AssertionError for 2 input channel image and label using Monai and pytorch

maryma: <p>Hi, I wanted to do the segmentation on kidneys dataset. For preprocessing data I tried to use cropforground and use labels so I could extract two kidneys. then for giving good dataset to training model I extract the left kidney and right kidney separably and I did these steps for labels. So I had left kidney images, left kidney labels, right kidney images and right kidney labels. I made two datasets with monai left kidney that contains left images and labels, right kidney that contains right images and labels. Because kidneys are two before making dataloader I tried to make a custom dataset with pytorch that gets left kidney dataset and right kidney dataset. Then concatenate the left image with right image and left label with right label. Finally it returns image and label which both of them has 2 channels. This is my custome dataset:</p>
<pre><code class="lang-auto">class Data_Single(Dataset):

    def __init__(self, im_dataset_l, im_dataset_r):
        self.im_dataset_l = im_dataset_l
        self.im_dataset_r = im_dataset_r

    def __len__(self):
        return len(self.im_dataset_l)

    def __getitem__(self, index):
        if type(index) is not int:
            if type(index) is not np.int64:
                raise ValueError(f"Need `index` to be `int`. Got {type(index)}.")
        if type(index) is np.int64:
            index = index.item()
            print(index)
                    
        img1 = self.im_dataset_l[index]['image']
        img2 = self.im_dataset_r[index]['image']
        
        
        lbl1 = self.im_dataset_l[index]['label']
        lbl2 = self.im_dataset_r[index]['label']
        
        img = torch.cat((img1,img2),0)
        lbl = torch.cat((lbl1,lbl2),0)
</code></pre>
<p>Then I chose Unet network that has 2 input channels and 2 output channels.</p>
<pre><code class="lang-auto">model = UNet(
    spatial_dims=3,
    in_channels=2,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(device)
</code></pre>
<p>The train loader is <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p><code>train_loader = DataLoader(Data_Single(im_dataset_l, im_dataset_r), batch_size = 10, shuffle=False)</code></p>
<p>I tried to extract image and label from batch data:</p>
<pre><code class="lang-auto">   for batch_data in train_loader:
        step += 1
        inputs, labels = (
            batch_data[0].to(device),
            batch_data[1].to(device)
                    )
</code></pre>
<p>And the output of batch data for image and label is:</p>
<pre><code class="lang-auto">torch.Size([1, 2, 128, 128, 128])
torch.Size([1, 2, 128, 128, 128])
</code></pre>
<p>But when I try to train model with this method:</p>
<pre><code class="lang-auto">outputs = model(inputs)
loss = loss_function(outputs, labels)
loss.backward()
</code></pre>
<p>I get this error on the second line of previous code:</p>
<p><code>raise AssertionError("labels should have a channel with length equal to one.")</code></p>
<p>I can not understand where is problem, can you please help me?</p>

ptrblck: <p>It seems the <code>loss_function</code> tries to call <code>one_hot</code> on the target tensor and fails <a href="https://github.com/Project-MONAI/MONAI/blob/d8d24eed0ee7705afd55050af0709613fddacb96/monai/networks/utils.py#L92">here</a>. Could you describe which loss function are are using as this seems to be a hard requirement?</p>

maryma: <p>My loss function was</p>
<p><code>loss_function = DiceLoss(to_onehot_y=True, softmax=True)</code></p>
<p>You are right the problem was onehot , when I changed it to false the problem was solved. But I couldn’t understand the logic of error I got that why did I get error when I have two channels and use onehot?</p>

ptrblck: <p>The <code>one_hot</code> operation would use class indices to create a one-hot encoded tensor from it as seen here:</p>
<pre><code class="lang-python">target = torch.randint(0, 2, (1, 4, 4))
print(target)
# tensor([[[1, 0, 0, 0],
#          [0, 1, 0, 1],
#          [0, 1, 1, 0],
#          [0, 1, 0, 0]]])

out = torch.nn.functional.one_hot(target, num_classes=2)
print(out)
# tensor([[[[0, 1],
#           [1, 0],
#           [1, 0],
#           [1, 0]],

#          [[1, 0],
#           [0, 1],
#           [1, 0],
#           [0, 1]],

#          [[1, 0],
#           [0, 1],
#           [0, 1],
#           [1, 0]],

#          [[1, 0],
#           [0, 1],
#           [1, 0],
#           [1, 0]]]])
</code></pre>
<p>I guess your two channels are already representing the classes and are one-hot encoded?</p>