Hello there! I am trying to train a model for segmentation for the first time. I found this way to load the model and the following loss function:
from torchvision.models.segmentation.deeplabv3 import DeepLabHead
def createDeepLabv3(outputchannels=2):
"""DeepLabv3 class with custom head
Args:
outputchannels (int, optional): The number of output channels
in your dataset masks. Defaults to 1.
Returns:
model: Returns the DeepLabv3 model with the ResNet101 backbone.
"""
model = models.segmentation.deeplabv3_resnet101(pretrained=True)
model.classifier = DeepLabHead(2048, outputchannels)
# Set the model in training mode
model.train()
return model
Loss function:
import torch.nn.functional as F
class DiceBCELoss(nn.Module):
def __init__(self, weight=None, size_average=True):
super(DiceBCELoss, self).__init__()
def forward(self, inputs, targets, smooth=1):
#comment out if your model contains a sigmoid or equivalent activation layer
inputs = F.sigmoid(inputs)
#flatten label and prediction tensors
print(inputs[0])
print(targets[0])
inputs = inputs.view(-1)
targets = targets.view(-1)
intersection = (inputs * targets).sum()
dice_loss = 1 - (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)
BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')
Dice_BCE = BCE + dice_loss
return Dice_BCE
I end up getting the following error:
AttributeError: âcollections.OrderedDictâ object has no attribute âsigmoidâ
AttributeError: âcollections.OrderedDictâ object has no attribute âviewâ
I assume that I loaded the model incorrectly.
As I understand it, the classifier of the pre-trained model is simply replaced here with a classifier with random weights. I already tried transfer learning and there I replaced only the last layer. So, I have a question, how to change the classifier correctly and how many classes should I specify, one or two, given that I am segmenting stones. And ofcource explain me please, how to load model correctly?