PyTorch量化模型训练时输出与目标尺寸不匹配问题求助

错误提示显示模型输出维度为[32,10]（对应10个分类），但one-hot编码后的标签维度仅为[32,9]，核心原因是标签的类别索引范围与你设定的num_classes不匹配，以下是针对性解决方法：

1. 强制指定one-hot编码的类别数 #

当前代码中torch.nn.functional.one_hot(target)会根据标签的最大索引自动生成维度，如果标签索引是1-10（而非0-9），会生成11维向量；如果标签仅包含0-8的索引，则生成9维。解决方法是手动指定num_classes参数，确保维度与模型输出一致：

target = torch.nn.functional.one_hot(target, num_classes=num_classes)

2. 验证标签的索引范围 #

先打印标签的最大/最小值，确认实际类别覆盖情况：

print("标签最小值:", target.min().item(), "标签最大值:", target.max().item())

• 如果最大值是9，说明类别索引是0-9，num_classes应设为10；
• 如果最大值是8，说明只有9个类别，需将模型最后一层的输出数改为9，同时修正num_classes值。

3. 检查损失函数是否适配 #

如果使用CrossEntropyLoss作为损失函数，不需要手动做one-hot编码，直接传入原始类别索引即可（该损失会自动处理索引到概率分布的转换），这能避免维度匹配问题：

# 去掉one-hot编码步骤
# target = torch.nn.functional.one_hot(target)
loss = criterion(output, target)  # 直接传入类别索引

修正后的训练函数示例 #

假设num_classes=10且标签索引正确，修改后的代码如下：

def train(model, train_loader, optimizer, criterion):
    losses = []
    model.train()    
    
    for i, data in enumerate(train_loader, 0):        
        inputs, target = data['pointcloud'].to(device).float(), data['category'].to(device)
        # 手动指定num_classes确保维度匹配
        target = torch.nn.functional.one_hot(target, num_classes=num_classes)
        optimizer.zero_grad()   
                
        output = model(inputs)
        loss = criterion(output, target.float())
        
        loss.backward()
        optimizer.step()
        
        losses.append(loss.data.cpu().numpy()) 
           
    return losses

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