张量维度不匹配

张量维度不匹配 | AI生成和翻译

Home 2025.04

看起来你在使用PyTorch和Transformers库进行人类反馈强化学习（RLHF）设置时，在训练奖励模型的过程中遇到了错误。错误信息表明MarginRankingLoss函数中存在维度不匹配，该函数期望所有输入张量（input1、input2和target）具有相同的大小。让我们来分析并解决这个问题。

错误分析

回溯显示：

RuntimeError: margin_ranking_loss : All input tensors should have same dimension but got sizes: input1: torch.Size([1, 1]), input2: torch.Size([1, 1]), target: torch.Size([1])

preferred_reward（input1）和less_preferred_reward（input2）的大小为[1, 1]
labels（target）的大小为[1]

MarginRankingLoss期望输入和目标具有兼容的维度。这里，preferred_reward和less_preferred_reward中的额外维度（即[1, 1] vs [1]）导致了不匹配。这很可能是因为奖励模型从线性层（self.v_head）输出带有额外维度的张量。

问题原因

在你的RewardModel中：

reward = self.v_head(last_hidden_states[:, -1])  # 形状: [batch_size, 1]

v_head层输出的奖励分数形状为[batch_size, 1]（例如，批量大小为1时为[1, 1]）。同时，labels被创建为：

labels = torch.ones(preferred_reward.size(0)).to(device)  # 形状: [batch_size]

这使得labels的形状为[1]，与奖励的[1, 1]形状不匹配。

解决方案

要解决这个问题，你需要确保奖励张量和目标张量具有兼容的形状。由于MarginRankingLoss期望1D张量（或至少匹配的形状），你可以从奖励输出中压缩额外的维度。以下是修改训练循环的方法：

更新后的代码片段

# 训练奖励模型
num_reward_epochs = 3
for epoch in range(num_reward_epochs):
    for prompt, preferred, less_preferred in preference_data:
        preferred_tokens = tokenizer(prompt + preferred, return_tensors="pt", truncation=True, max_length=128).to(device)
        less_preferred_tokens = tokenizer(prompt + less_preferred, return_tensors="pt", truncation=True, max_length=128).to(device)

        preferred_reward = reward_model(**preferred_tokens).squeeze()  # 移除额外维度
        less_preferred_reward = reward_model(**less_preferred_tokens).squeeze()  # 移除额外维度

        labels = torch.ones(preferred_reward.size(0)).to(device)  # 形状: [batch_size]，例如[1]
        loss = reward_criterion(preferred_reward, less_preferred_reward, labels)

        reward_optimizer.zero_grad()
        loss.backward()
        reward_optimizer.step()
    print(f"奖励轮次 {epoch+1}, 损失: {loss.item()}")

关键更改

在preferred_reward和less_preferred_reward上添加.squeeze()，将它们的形状从[1, 1]减少到[1]，与labels张量匹配。

补充说明

批量大小考虑：
- 如果之后增加批量大小（例如一次处理多个样本），请确保正确批处理preferred_tokens和less_preferred_tokens，并且在压缩后形状仍然对齐。
奖励模型输出：
- 当前的奖励模型获取最后一个隐藏状态（last_hidden_states[:, -1]）并通过线性层传递。对于简单的奖励模型来说这没问题，但你可以尝试聚合隐藏状态（例如平均池化）以获得更细致的评分。
RL循环：
- 你的RL微调循环使用了简化的损失（-reward.mean()）。实际上，RLHF通常采用像PPO（近端策略优化）这样的算法，需要更复杂的设置（例如价值函数、裁剪目标）。目前你的方法作为概念验证是可行的。
分词器填充：
- 你可能需要在分词器调用中添加padding=True以确保输入长度一致，不过带有max_length的truncation=True在这里应该能处理大多数情况。

完整工作示例（简化版）

以下是修正后的奖励训练循环的简化版本：

import torch
import torch.nn as nn
import torch.optim as optim
from transformers import AutoModelForCausalLM, AutoTokenizer

# 示例数据
preference_data = [
    ("写一个关于猫的故事。", "Whiskers the cat lived in a cozy cottage...", "一个关于猫的故事。结束。"),
]

# 加载模型和分词器
model_name = "gpt2"
policy_model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
device = "cuda" if torch.cuda.is_available() else "cpu"
policy_model.to(device)

# 奖励模型
class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.base_model = base_model.transformer
        self.v_head = nn.Linear(base_model.config.n_embd, 1)

    def forward(self, input_ids, attention_mask):
        outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)
        last_hidden_states = outputs.last_hidden_state
        reward = self.v_head(last_hidden_states[:, -1])
        return reward

reward_model = RewardModel(policy_model).to(device)
reward_optimizer = optim.AdamW(reward_model.parameters(), lr=1e-5)
reward_criterion = nn.MarginRankingLoss(margin=1.0)

# 训练奖励模型
for epoch in range(3):
    for prompt, preferred, less_preferred in preference_data:
        preferred_tokens = tokenizer(prompt + preferred, return_tensors="pt", truncation=True, max_length=128).to(device)
        less_preferred_tokens = tokenizer(prompt + less_preferred, return_tensors="pt", truncation=True, max_length=128).to(device)

        preferred_reward = reward_model(**preferred_tokens).squeeze()
        less_preferred_reward = reward_model(**less_preferred_tokens).squeeze()

        labels = torch.ones(preferred_reward.size(0)).to(device)
        loss = reward_criterion(preferred_reward, less_preferred_reward, labels)

        reward_optimizer.zero_grad()
        loss.backward()
        reward_optimizer.step()
    print(f"奖励轮次 {epoch+1}, 损失: {loss.item()}")

这应该能解决维度不匹配问题，让你的奖励模型成功训练。如果你需要进一步澄清或帮助RL部分，请告诉我！

理解和修复PyTorch脚本中的IndexError

你在Python脚本中遇到的错误信息是：

IndexError: Dimension specified as 0 but tensor has no dimensions

这发生在你的脚本rl.py的第44行，具体是以下这行代码：

labels = torch.ones(preferred_reward.size(0)).to(device)

这个错误表明preferred_reward是一个标量张量（即它没有维度，形状为[]），尝试用.size(0)访问其第一个维度是无效的，因为标量没有可以索引的维度。

让我们逐步分析这个问题并提供解决方案。

错误原因

在你的脚本中，你很可能正在使用像MarginRankingLoss这样的损失函数来训练奖励模型，该函数期望其输入（preferred_reward、less_preferred_reward和labels）是具有兼容形状的张量——通常是1D张量，其中每个元素对应批次中的一个样本。以下是正在发生的情况：

preferred_reward的来源：
- preferred_reward是reward_model前向传播的输出，例如reward_model(**preferred_tokens)。
- 假设你的奖励模型为每个样本输出单个值，对于批量大小为1的情况，输出形状为[1, 1]（批量大小 × 输出维度）。
压缩张量：
- 在你的原始代码中，你对preferred_reward应用了.squeeze()：
```
preferred_reward = reward_model(**preferred_tokens).squeeze()
```
- .squeeze()方法移除所有大小为1的维度。对于形状为[1, 1]的张量，这会将其减少为[]——一个没有维度的标量张量。
访问大小：
- 随后，你尝试创建一个与preferred_reward具有相同批量大小的labels张量：
```
labels = torch.ones(preferred_reward.size(0)).to(device)
```
- 对于标量张量（[]），preferred_reward.size()返回torch.Size([])，一个空的大小元组。尝试用.size(0)访问第一个维度会引发IndexError，因为没有维度可访问。
预期行为：
- MarginRankingLoss函数要求其输入（preferred_reward、less_preferred_reward和labels）具有相同的形状，通常是像[batch_size]这样的1D张量。对于批量大小为1的情况，这应该是[1]，而不是标量[]。

根本原因是.squeeze()过于激进——它移除了所有单例维度，当批量大小为1时将[1, 1]变成了标量[]，这破坏了后续的.size(0)调用。

解决方案

要解决这个问题，你需要确保preferred_reward和less_preferred_reward保持为形状为[batch_size]的1D张量，即使在batch_size=1时也是如此。不要使用会移除所有大小为1的维度的.squeeze()，而是使用.squeeze(-1)仅移除最后一个维度。这将把[1, 1]转换为[1]，保留批次维度。

以下是修正后的奖励模型训练循环代码片段：

# 训练奖励模型
num_reward_epochs = 3
for epoch in range(num_reward_epochs):
    for prompt, preferred, less_preferred in preference_data:
        # 分词输入
        preferred_tokens = tokenizer(prompt + preferred, return_tensors="pt", truncation=True, max_length=128).to(device)
        less_preferred_tokens = tokenizer(prompt + less_preferred, return_tensors="pt", truncation=True, max_length=128).to(device)

        # 计算奖励，仅压缩最后一个维度
        preferred_reward = reward_model(**preferred_tokens).squeeze(-1)  # 形状: [1]
        less_preferred_reward = reward_model(**less_preferred_tokens).squeeze(-1)  # 形状: [1]

        # 基于批量大小创建标签张量
        labels = torch.ones(preferred_reward.size(0)).to(device)  # 形状: [1]
        
        # 计算损失
        loss = reward_criterion(preferred_reward, less_preferred_reward, labels)

        # 反向传播
        reward_optimizer.zero_grad()
        loss.backward()
        reward_optimizer.step()
    
    print(f"奖励轮次 {epoch+1}, 损失: {loss.item()}")

工作原理

在.squeeze(-1)之后：
- 来自reward_model的原始形状：[1, 1]（batch_size=1, output_dim=1）
- 在.squeeze(-1)之后：[1]（仅移除最后一个维度）
- preferred_reward.size(0)返回1，即批量大小
- labels变为形状为[1]的1D张量，与preferred_reward和less_preferred_reward的形状匹配
与MarginRankingLoss的兼容性：
- MarginRankingLoss期望input1（preferred_reward）、input2（less_preferred_reward）和target（labels）具有相同的形状。当三者都是[1]时，损失计算可以顺利进行而不会出错。
可扩展性：
- 如果你之后使用更大的批量大小（例如2），奖励模型输出[2, 1]，.squeeze(-1)产生[2]，而labels变为[2]，保持一致性。

替代方法

虽然.squeeze(-1)是一个干净且精确的修复方法，但以下两种方法也同样有效：

使用索引：

preferred_reward = reward_model(**preferred_tokens)[:, 0]  # 形状: [1]
less_preferred_reward = reward_model(**less_preferred_tokens)[:, 0]  # 形状: [1]

这选择了最后一个维度的第一个（也是唯一的）元素，将[1, 1]转换为[1]。

使用.view(-1)：

preferred_reward = reward_model(**preferred_tokens).view(-1)  # 形状: [1]
less_preferred_reward = reward_model(**less_preferred_tokens).view(-1)  # 形状: [1]

这将张量展平为1D张量。对于[1, 1]，由于只有一个元素，它变为[1]。

这两种替代方法在此上下文中都能达到与.squeeze(-1)相同的结果，但.squeeze(-1)是首选，因为它明确针对最后一个维度，与奖励模型的输出结构（[batch_size, 1]）对齐。

补充说明

输出中的警告：
- 关于torch.utils._pytree._register_pytree_node和resume_download的FutureWarning消息与IndexError无关。它们来自transformers库，表明将在未来版本中更改的已弃用函数或行为。目前，它们不会影响你的脚本功能，除非你希望更新代码以获得未来兼容性。
奖励模型输出：
- 该修复假设你的reward_model输出形状为[batch_size, 1]的张量。如果其架构不同（例如直接输出标量或不同形状），你需要相应调整。基于典型的奖励模型设计，[batch_size, 1]在线性头之后是标准的。
批处理：
- 你当前的循环一次处理一个偏好对（隐式批量大小为1）。如果你计划扩展规模，考虑使用DataLoader批处理多个对以提高效率，但.squeeze(-1)修复仍然适用。

验证

对于批量大小为1：

reward_model(**preferred_tokens)：[1, 1]
在.squeeze(-1)之后：[1]
preferred_reward.size(0)：1
labels：[1]
损失输入：全部为[1]——完美运行。

IndexError被消除，你的奖励模型训练现在应该能顺利运行。

最终答案

要修复rl.py第44行的IndexError: Dimension specified as 0 but tensor has no dimensions，在计算preferred_reward和less_preferred_reward时将.squeeze()替换为.squeeze(-1)。这确保它们保持形状为[1]（或对于更大的批次为[batch_size]）的1D张量，允许preferred_reward.size(0)正确检索批量大小以创建labels张量。

按照上面所示在你的奖励模型训练循环中应用此更改，错误将被解决。

Back Donate