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

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:

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)

Then I chose Unet network that has 2 input channels and 2 output channels.

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)

The train loader is :slight_smile:

train_loader = DataLoader(Data_Single(im_dataset_l, im_dataset_r), batch_size = 10, shuffle=False)

I tried to extract image and label from batch data:

   for batch_data in train_loader:
        step += 1
        inputs, labels = (
            batch_data[0].to(device),
            batch_data[1].to(device)
                    )

And the output of batch data for image and label is:

torch.Size([1, 2, 128, 128, 128])
torch.Size([1, 2, 128, 128, 128])

But when I try to train model with this method:

outputs = model(inputs)
loss = loss_function(outputs, labels)
loss.backward()

I get this error on the second line of previous code:

raise AssertionError("labels should have a channel with length equal to one.")

I can not understand where is problem, can you please help me?

It seems the loss_function tries to call one_hot on the target tensor and fails here. Could you describe which loss function are are using as this seems to be a hard requirement?

1 Like

My loss function was

loss_function = DiceLoss(to_onehot_y=True, softmax=True)

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?

The one_hot operation would use class indices to create a one-hot encoded tensor from it as seen here:

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]]]])

I guess your two channels are already representing the classes and are one-hot encoded?

1 Like