输入空特征时SparseAdam引发维度不匹配RuntimeError求助

运行环境 #

• PyTorch版本：1.8.2
• 运行模式：CPU

复现代码 #

from torch import nn,optim
import torch
import torch.nn.functional as F
from torch.autograd import Variable

def cumsum_1d_with_zero(input: torch.Tensor, include_last: bool = False) -> torch.Tensor:
    output = torch.zeros(input.size(0) + 1, dtype=input.dtype)
    torch.cumsum(input, 0, out=output[1:])
    if include_last:
        return output
    else:
        return output[:-1]

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.embedding = nn.EmbeddingBag(4, 300, sparse=True, mode='sum')
        self.linear1 = nn.Linear(300, 1)
        self.optimizer = optim.SparseAdam(self.embedding.parameters(), 0.0001)
        self.dense_optimizer = optim.Adam(self.linear1.parameters(), 0.0001)
    def forward(self,inputs):
        rows = [torch.tensor(input, dtype=torch.int64) for input in inputs]
        lengths = torch.tensor([x.size(0) for x in rows], dtype=torch.int64)
        offsets = cumsum_1d_with_zero(lengths)
        indices = torch.cat(rows)
        embedding_res = self.embedding(indices, offsets)
        print(embedding_res)
        logits = F.sigmoid(self.linear1(embedding_res))
        return logits

    def optimize(self, logits):
        labels = Variable(torch.LongTensor([0, 1, 0, 1]).reshape(4,1).float())
        criterion = nn.BCEWithLogitsLoss()
        loss = criterion(logits, labels)
        print(loss.item())
        loss.backward()
        self.optimizer.step()
        self.dense_optimizer.step()

model = MyModel()
inputs = [[],[],[],[]]
result = model(inputs)
model.optimize(result)

报错信息 #

RuntimeError: The size of tensor a (300) must match the size of tensor b (0) at non-singleton dimension 1

核心原因 #

错误源于**EmbeddingBag的offsets参数格式不符合要求**，导致后续线性层计算时出现维度不匹配：

1. EmbeddingBag要求offsets是长度为num_bags + 1的1D张量（num_bags是输入样本的数量，这里为4），其中offsets[i]表示第i个bag在indices中的起始索引，offsets[-1]是indices的总长度。
2. 代码中cumsum_1d_with_zero函数使用了默认参数include_last=False，生成的offsets长度为4（与num_bags相等），违反了num_bags + 1的参数要求。
3. 当输入全为空列表时，indices是空张量，配合长度错误的offsets，EmbeddingBag输出的张量特征维度异常（变为0），而线性层linear1的输入特征维度要求是300，矩阵乘法时就会出现"300 vs 0"的维度不匹配错误。

修复方案 #

将forward方法中生成offsets的代码改为：

offsets = cumsum_1d_with_zero(lengths, include_last=True)

这样生成的offsets长度为5（4+1），符合EmbeddingBag的参数要求。此时空bag会被正确处理为全0的300维向量，embedding_res形状为(4, 300)，后续线性层和loss计算都能正常执行。

内容的提问来源于stack exchange，提问作者Zhu Kai hua