创建PyTorch分割模型时遇AttributeError：无'to'属性问题求助

错误原因 #

你的segmentationmodels类没有继承PyTorch核心的nn.Module基类，因此无法调用to()这类PyTorch模型专属的方法，导致抛出AttributeError。

修复代码 #

只需让你的类继承nn.Module即可，修改后的完整代码如下：

from torch import nn
import segmentation_models_pytorch as smp
from segmentation_models_pytorch.losses import DiceLoss 

# 继承nn.Module基类
class segmentationmodels(nn.Module):
  def __init__(self):
    super(segmentationmodels,self).__init__()
    # 需提前定义ENCODER和weights变量，例如ENCODER='resnet34'，weights='imagenet'
    self.arc=smp.Unet(
        encoder_name=ENCODER,
        encoder_weights=weights,
        in_channels=3,
        classes=1,
        activation=None
    )
  def forward(self,images,mask=None):
    logits=self.arc(images)
    if mask is not None:  # 推荐用is not None替代!=None，符合Python规范
      loss1=DiceLoss(mode='binary')(logits,mask)
      loss2=nn.BCEWithLogitsLoss()(logits,mask)
      return logits,loss1+loss2
    return logits

# 需提前定义DEVICE，例如DEVICE='cuda' if torch.cuda.is_available() else 'cpu'
model=segmentationmodels()
model.to(DEVICE)

额外注意事项 #

• 必须提前定义ENCODER、weights和DEVICE变量，否则会触发新的未定义错误。
• 判断变量非空时，mask is not None是更规范的Python写法，建议替换原代码中的mask != None。

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