0

This is my code:

from torch import nn
import segmentation_models_pytorch as smp
from segmentation_models_pytorch.losses import DiceLoss 
class segmentationmodels():
  def __init__(self):
    super(segmentationmodels,self).__init__()
    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 != None:
      loss1=DiceLoss(mode='binary')(logits,mask)
      loss2=nn.BCEWithLogitsLoss()(logits,mask)
      return logits,loss1+loss2
    return logits
model=segmentationmodels()
model.to(DEVICE)

And this is the error:

AttributeError                            Traceback (most recent call last)
<ipython-input-123-b61ad7737aab> in <module>
      1 model=segmentationmodels()
----> 2 model.to(DEVICE)

AttributeError: 'segmentationmodels' object has no attribute 'to`
desertnaut
  • 57,590
  • 26
  • 140
  • 166

1 Answers1

0

You didn't mention torch.nn.Module when creating the class, that's why you are getting that error. So, your code should be this:

from torch import nn
import segmentation_models_pytorch as smp
from segmentation_models_pytorch.losses import DiceLoss 
class segmentationmodels(torch.nn.Module):
  def __init__(self):
    super(segmentationmodels,self).__init__()
    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 != None:
      loss1=DiceLoss(mode='binary')(logits,mask)
      loss2=nn.BCEWithLogitsLoss()(logits,mask)
      return logits,loss1+loss2
    return logits
model=segmentationmodels()
model.to(DEVICE)
MD Mushfirat Mohaimin
  • 1,966
  • 3
  • 10
  • 22