解决预训练RoBERTa信号生成中的张量维度不匹配RuntimeError

我尝试用预训练RoBERTa模型学习已有信号并生成新信号，但训练时触发错误：RuntimeError: The size of tensor a (768) must match the size of tensor b (374125) at non-singleton dimension 1，我的实现代码如下，求改进方案：

import torch
from transformers import RobertaTokenizer, RobertaModel
from torch.nn import Linear
from torch.optim import Adam

# Load RoBERTa tokenizer and model
tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
model = RobertaModel.from_pretrained('roberta-base')

# Add a linear layer on top of RoBERTa to generate signal
signal_size = 2993 * 125
model.classifier = Linear(model.config.hidden_size, 374125)
# Prepare your signal data as input sequences and target signals
input_sequences = [tokenizer.encode("signal {}".format(i), return_tensors='pt') for i in range(16)]
target_signals = [train_all_data.transpose(1,3,0,2).reshape(16, signal_size)[i] for i in range(16)]
target_signals = [torch.tensor(signal) for signal in target_signals]

# Define loss function and optimizer
criterion = torch.nn.MSELoss()
optimizer = Adam(model.parameters(), lr=0.001)

# Train the model on your signal data
for epoch in range(1000):
    for input_ids, target in zip(input_sequences, target_signals):
        optimizer.zero_grad()
        outputs = model(input_ids)[0].mean(dim=1)
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()
        
# Use the fine-tuned model to generate a new signal
input_ids = torch.tensor(tokenizer.encode("Generate a new signal", return_tensors='pt')).unsqueeze(0)
outputs = model(input_ids)[0].mean(dim=1)
new_signal = outputs.detach().numpy()

问题根源 #

错误核心是模型输出与目标张量维度不匹配：

• RoBERTa池化后输出的张量维度是(1, 768)（768为RoBERTa-base的hidden_size）
• 目标信号的维度是(374125)，两者无法直接计算MSE损失

修正方案 #

1. 正确使用分类头 #

你已经定义了映射到目标维度的线性层model.classifier，但训练时没有调用该层，直接使用了RoBERTa的原始池化输出。需添加分类头的前向传播：

outputs = model(input_ids)[0].mean(dim=1)
# 新增：通过分类头将768维映射到374125维
outputs = model.classifier(outputs)

2. 统一张量维度 #

目标信号需要保持与模型输出一致的batch维度（即使单样本，也需为(1, 374125)而非(374125)），处理目标时添加维度：

target_signals = [torch.tensor(signal).unsqueeze(0) for signal in target_signals]

3. 优化数据处理（可选） #

避免用列表循环处理样本，直接整合成batch张量提升训练效率：

# 合并输入序列为单个batch张量
input_ids = torch.cat([tokenizer.encode("signal {}".format(i), return_tensors='pt') for i in range(16)], dim=0)
# 合并目标信号为单个batch张量
target_signals = torch.tensor(train_all_data.transpose(1,3,0,2).reshape(16, signal_size))

4. 修正信号生成逻辑 #

生成新信号时同样需要通过分类头完成维度映射：

input_ids = tokenizer.encode("Generate a new signal", return_tensors='pt')
outputs = model(input_ids)[0].mean(dim=1)
outputs = model.classifier(outputs)  # 新增分类头调用
new_signal = outputs.detach().numpy()

完整修正后的代码 #

import torch
from transformers import RobertaTokenizer, RobertaModel
from torch.nn import Linear
from torch.optim import Adam

# Load RoBERTa tokenizer and model
tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
model = RobertaModel.from_pretrained('roberta-base')

# Add a linear layer on top of RoBERTa to generate signal
signal_size = 2993 * 125
model.classifier = Linear(model.config.hidden_size, signal_size)

# Prepare your signal data as input sequences and target signals
input_ids = torch.cat([tokenizer.encode("signal {}".format(i), return_tensors='pt') for i in range(16)], dim=0)
target_signals = torch.tensor(train_all_data.transpose(1,3,0,2).reshape(16, signal_size))

# Define loss function and optimizer
criterion = torch.nn.MSELoss()
optimizer = Adam(model.parameters(), lr=0.001)

# Train the model on your signal data
for epoch in range(1000):
    optimizer.zero_grad()
    roberta_outputs = model(input_ids)[0].mean(dim=1)
    outputs = model.classifier(roberta_outputs)
    loss = criterion(outputs, target_signals)
    loss.backward()
    optimizer.step()
    if epoch % 50 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item()}")
        
# Use the fine-tuned model to generate a new signal
input_ids = tokenizer.encode("Generate a new signal", return_tensors='pt')
roberta_outputs = model(input_ids)[0].mean(dim=1)
outputs = model.classifier(roberta_outputs)
new_signal = outputs.detach().numpy()

内容的提问来源于stack exchange，提问作者assa