Ubuntu环境A100运行CNN代码遇张量类型不匹配RuntimeError

核心原因 #

模型参数是torch.cuda.FloatTensor（默认单精度浮点），但输入数据是torch.cuda.DoubleTensor（双精度），前向传播时类型不兼容导致报错。Windows和Linux环境下数据加载/生成的默认类型存在差异，才会出现之前正常、现在报错的情况。

具体解决步骤 #

1. 统一模型与输入的类型 #

• 若要让模型适配输入的DoubleTensor：

  model = model.double()
  

• 更推荐让输入适配模型的FloatTensor（GPU训练单精度足够，速度更快、显存占用更低）：

  # 输入模型前，把所有输入张量转成float类型
  x = x.float()
  

  注意：必须保证所有输入数据流（包括标签、中间计算变量）和模型参数类型完全一致，只转部分变量没用。

2. 从数据加载环节根治问题 #

很多时候问题出在数据加载时默认生成了双精度数据，直接在加载阶段指定类型：

• 用numpy读取数据时：

  data = np.load("your_data_path.npy").astype(np.float32)
  

• 自定义Dataset类时，在__getitem__里明确转换：

  def __getitem__(self, idx):
      image = self.data[idx]
      image = torch.tensor(image, dtype=torch.float32)
      label = torch.tensor(self.labels[idx], dtype=torch.long)
      return image, label
  

3. 检查损失函数的类型 #

损失函数会继承输入的类型，如果输入是DoubleTensor，要同步转换损失函数：

criterion = nn.CrossEntropyLoss().double()

当然更省事的是统一用FloatTensor，不用额外调整损失函数。

4. 调试定位类型不匹配的位置 #

可以在模型前向传播前加打印语句，快速确认哪部分类型不对：

print("输入张量类型：", x.dtype)
print("模型第一层权重类型：", model.conv1.weight.dtype)

额外建议 #

A100对单精度（float32）的优化远好于双精度（float64），统一用float32不仅能解决问题，还能大幅提升训练速度、减少显存消耗。

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