PyTorch图像分割模型预测报错：'NoneType' object has no attribute 'size'

问题背景 #

无GPU、未使用CUDA，在CPU上基于PyTorch训练图像分割模型后，执行预测代码时触发错误：

AttributeError: 'NoneType' object has no attribute 'size'

执行的预测代码 #

idx = 20
model.load_state_dict(torch.load('/content/best_model.pt'))

image, mask = validset[idx]
image = image.unsqueeze_(0)
print(type(image))
# logits_mask = model(image.to(DEVICE).unsqueeze(0)) # (c,h,w) -> (1,c,h,w)
logits_mask = model(image) # (c,h,w) py-> (1,c,h,w)

完整报错堆栈 #

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-56-edf3f0fae49c> in <module>
      6 print(type(image))
      7 # logits_mask = model(image.to(DEVICE).unsqueeze(0)) # (c,h,w) -> (1,c,h,w)
----> 8 logits_mask = model(image) # (c,h,w) py-> (1,c,h,w)
      9 
     10 pred_mask = torch.sigmoid(logits_mask)

3 frames
/usr/local/lib/python3.7/dist-packages/segmentation_models_pytorch/losses/dice.py in forward(self, y_pred, y_true)
     57     def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
     58 
---> 59         assert y_true.size(0) == y_pred.size(0)
     60 
     61         if self.from_logits:

AttributeError: 'NoneType' object has no attribute 'size'

问题分析与解决 #

错误根源是segmentation_models_pytorch的DiceLoss在执行时，y_true参数为None。这说明你的模型forward流程被错误绑定了损失计算逻辑，导致预测时仅传入图像输入，却触发了需要标签的损失计算步骤。

具体修复步骤： #

1. 分离模型与损失函数
   检查模型定义，确保forward方法仅负责处理输入图像并返回预测logits，不要将损失计算写在forward内部。训练时应在外部单独计算损失：

   # 训练时的正确流程示例
   logits = model(train_image)
   loss = dice_loss(logits, train_mask)
   

2. 预测前切换模型到eval模式
   加载模型后必须执行model.eval()，避免BatchNorm、Dropout等层处于训练模式，同时防止意外触发训练相关逻辑：

   model.load_state_dict(torch.load('/content/best_model.pt'))
   model.eval()  # 新增该行
   

3. 确保模型与输入设备一致
   虽然使用CPU，但需确认加载的模型权重处于CPU设备，可手动指定：

   model.load_state_dict(torch.load('/content/best_model.pt', map_location=torch.device('cpu')))
   model.to('cpu')
   

4. 验证输入维度
   确认输入张量维度为(batch_size, channels, height, width)，你的代码中image.unsqueeze_(0)已添加batch维度，可打印确认：

   print(image.shape)  # 预期输出为 (1, C, H, W)
   

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