Im using a Pytorch Unet model to which i am feeding in a image as input and along with that i am feeding the label as the input image mask and traning the dataset on it. The Unet model i have picked up from somewhere else, and i am using the cross-entropy loss as a loss function but i get this dimension out of range error,
RuntimeError Traceback (most recent call last) <ipython-input-358-fa0ef49a43ae> in <module>() 16 for epoch in range(0, num_epochs): 17 # train for one epoch ---> 18 curr_loss = train(train_loader, model, criterion, epoch, num_epochs) 19 20 # store best loss and save a model checkpoint <ipython-input-356-1bd6c6c281fb> in train(train_loader, model, criterion, epoch, num_epochs) 16 # measure loss 17 print (outputs.size(),labels.size()) ---> 18 loss = criterion(outputs, labels) 19 losses.update(loss.data[0], images.size(0)) 20 /usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py in _ _call__(self, *input, **kwargs) 323 for hook in self._forward_pre_hooks.values(): 324 hook(self, input) --> 325 result = self.forward(*input, **kwargs) 326 for hook in self._forward_hooks.values(): 327 hook_result = hook(self, input, result) <ipython-input-355-db66abcdb074> in forward(self, logits, targets) 9 probs_flat = probs.view(-1) 10 targets_flat = targets.view(-1) ---> 11 return self.crossEntropy_loss(probs_flat, targets_flat) /usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs) 323 for hook in self._forward_pre_hooks.values(): 324 hook(self, input) --> 325 result = self.forward(*input, **kwargs) 326 for hook in self._forward_hooks.values(): 327 hook_result = hook(self, input, result) /usr/local/lib/python3.5/dist-packages/torch/nn/modules/loss.py in f orward(self, input, target) 599 _assert_no_grad(target) 600 return F.cross_entropy(input, target, self.weight, self.size_average, --> 601 self.ignore_index, self.reduce) 602 603 /usr/local/lib/python3.5/dist-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce) 1138 >>> loss.backward() 1139 """ -> 1140 return nll_loss(log_softmax(input, 1), target, weight, size_average, ignore_index, reduce) 1141 1142 /usr/local/lib/python3.5/dist-packages/torch/nn/functional.py in log_softmax(input, dim, _stacklevel) 784 if dim is None: 785 dim = _get_softmax_dim('log_softmax', input.dim(), _stacklevel) --> 786 return torch._C._nn.log_softmax(input, dim) 787 788 RuntimeError: dimension out of range (expected to be in range of [-1, 0], but got 1)
Part of my code looks like this
class crossEntropy(nn.Module): def __init__(self, weight = None, size_average = True): super(crossEntropy, self).__init__() self.crossEntropy_loss = nn.CrossEntropyLoss(weight, size_average) def forward(self, logits, targets): probs = F.sigmoid(logits) probs_flat = probs.view(-1) targets_flat = targets.view(-1) return self.crossEntropy_loss(probs_flat, targets_flat) class UNet(nn.Module): def __init__(self, imsize): super(UNet, self).__init__() self.imsize = imsize self.activation = F.relu self.pool1 = nn.MaxPool2d(2) self.pool2 = nn.MaxPool2d(2) self.pool3 = nn.MaxPool2d(2) self.pool4 = nn.MaxPool2d(2) self.conv_block1_64 = UNetConvBlock(4, 64) self.conv_block64_128 = UNetConvBlock(64, 128) self.conv_block128_256 = UNetConvBlock(128, 256) self.conv_block256_512 = UNetConvBlock(256, 512) self.conv_block512_1024 = UNetConvBlock(512, 1024) self.up_block1024_512 = UNetUpBlock(1024, 512) self.up_block512_256 = UNetUpBlock(512, 256) self.up_block256_128 = UNetUpBlock(256, 128) self.up_block128_64 = UNetUpBlock(128, 64) self.last = nn.Conv2d(64, 2, 1)