训练Unet模型时出现NotImplementedError报错如何解决

报错根因 #

这个报错是torch.nn.Module基类的默认兜底逻辑抛出的：当你调用model(输入)时，Python找不到模型类里定义的forward()方法，就会fallback到父类的空forward，直接抛出未实现错误，和你的训练/验证循环逻辑本身无关，问题出在Unet模型定义或者实例化环节。

排查修复步骤 #

• 第一步：检查Unet类的forward方法定义
  方法名必须严格拼写为forward，不能出现foward/forword/Forward这类大小写、拼写错误，否则Python无法识别这是模型的前向传播方法。
  同时注意你训练时调用模型传入了images、masks两个参数，forward方法的入参需要和调用逻辑匹配，参考正确写法：

  import torch.nn as nn
  class Unet(nn.Module):
      def __init__(self):
          super().__init__()
          # 此处填写你的编码器、解码器、卷积层定义
      
      # 方法名必须是forward，不能拼错
      def forward(self, images, masks=None):
          # 填写前向计算逻辑，输出预测logits
          logits =  # 卷积层堆叠计算结果
          if masks is not None:
              # 计算损失
              loss = nn.BCEWithLogitsLoss()(logits, masks) # 替换为你实际用的损失函数
              return logits, loss
          return logits
  

• 第二步：检查模型实例化逻辑
  不能直接把Unet类本身传入训练函数，必须加括号完成实例化，错误写法和正确写法对比如下：

  # 错误：把类对象当成实例传入，调用时会触发报错
  # model = Unet
  # 正确：实例化后再迁移到对应设备
  model = Unet() # 如有初始化参数按需传入
  model = model.to(DEVICE)
  

  如果你的Unet是拆分了下采样、上采样等多个自定义子模块拼接而成，也要逐个检查子模块的forward方法有没有拼写错误。
• 第三步：提前做前向测试快速定位问题
  不需要启动完整训练循环，在训练代码前加几行测试逻辑，用随机张量验证前向传播是否通畅：

  # 张量尺寸和你实际数据集的输入尺寸对齐即可
  test_imgs = torch.randn(2, 3, 256, 256).to(DEVICE)
  test_masks = torch.randn(2, 1, 256, 256).to(DEVICE)
  logits, loss = model(test_imgs, test_masks)
  print(f"前向测试通过，输出维度：{logits.shape}, 测试损失：{loss.item()}")
  

  如果这段代码能正常运行不报错，说明模型定义没问题，再检查训练循环的缩进问题即可。

附带修复训练循环缩进问题 #

你贴的训练循环中，epoch循环内的代码没有正确缩进，Python会把循环外的代码当成循环体，运行时也会触发逻辑错误，修正后写法：

best_valid_loss = float('inf') # 不依赖numpy也可以直接用float('inf')
for i in range(EPOCHS):
    # 循环内代码统一缩进4个空格
    train_loss = train_fn(trainloader, model, optimizer)
    valid_loss = eval_fn(validloader, model)

    if valid_loss < best_valid_loss:
        torch.save(model.state_dict(), 'best_model.pt')
        print("SAVED_MODEL")
        best_valid_loss = valid_loss
    
    print(f"Epoch : {i+1} Train_loss: {train_loss} Valid_loss: {valid_loss}")

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