Hi!
I created my own U-Net model following some walkthrough tutorials, trying to be as close to original paper as possible, making some small changes as nb_classes, padding. Model looks like this:
def double_conv(in_c, out_c):
    conv = nn.Sequential(
        nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_c, out_c, kernel_size=3, padding=1),
        nn.ReLU(inplace=True)
    )
    return conv

def crop_img(tensor, target_tensor):
    # square images
    target_size = target_tensor.size()[2]
    if tensor.size()[2] % 2 == 1:
        tensor_size = tensor.size()[2]-1
    else:
        tensor_size = tensor.size()[2]
    delta = tensor_size - target_size
    delta = delta // 2
    return tensor[:, :, delta:tensor_size-delta, delta:tensor_size-delta]

class UNet(nn.Module):
    def __init__(self, nb_classes):
        super(UNet, self).__init__()
        self.max_pool_2x2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.down_conv_1 = double_conv(3, 64)
        self.down_conv_2 = double_conv(64, 128)
        self.down_conv_3 = double_conv(128, 256)
        self.down_conv_4 = double_conv(256, 512)
        self.down_conv_5 = double_conv(512, 1024)

        ## transposed convolutions
        self.up_trans_1 = nn.ConvTranspose2d(1024, 512, 2, 2)
        self.up_conv_1 = double_conv(1024, 512)

        self.up_trans_2 = nn.ConvTranspose2d(512, 256, 2, 2)
        self.up_conv_2 = double_conv(512, 256)

        self.up_trans_3 = nn.ConvTranspose2d(256, 128, 2, 2)
        self.up_conv_3 = double_conv(256, 128)

        self.up_trans_4 = nn.ConvTranspose2d(128, 64, 2, 2)
        self.up_conv_4 = double_conv(128, 64)

        self.out = nn.Conv2d(64, nb_classes, 1)


    def forward(self, image):
        # encoder part
        # input image
        x1 = self.down_conv_1(image) # this is passed to decoder
        # max pooling
        x2 = self.max_pool_2x2(x1)

        x3 = self.down_conv_2(x2) # this is passed to decoder
        x4 = self.max_pool_2x2(x3)

        x5 = self.down_conv_3(x4) # this is passed to decoder
        x6 = self.max_pool_2x2(x5)

        x7 = self.down_conv_4(x6) # this is passed to decoder
        x8 = self.max_pool_2x2(x7)

        x9 = self.down_conv_5(x8)

        # decoder part
        x = self.up_trans_1(x9)
        y = crop_img(x7, x)
        x = self.up_conv_1(torch.cat([x, y], 1))

        x = self.up_trans_2(x)
        y = crop_img(x5, x)
        x = self.up_conv_2(torch.cat([x, y], 1))

        x = self.up_trans_3(x)
        y = crop_img(x3, x)
        x = self.up_conv_3(torch.cat([x, y], 1))

        x = self.up_trans_4(x)
        y = crop_img(x1, x)
        x = self.up_conv_4(torch.cat([x, y], 1))

        x = self.out(x)
        return x

Then before jumping to some larger data, I wanted to learn and understand how everything works, so I created images and masks using this github.
image.shape:   (512, 512, 3)
image.min(), image.max():   0.0 255.0
mask.shape:   (6, 512, 512)
mask.min(), mask.max():   0.0 1.0

So the masks are 6-channel arrays, where each channel represents one color and sometimes channels overlay themselves, as on this example image below. 
image_0_mask512×512 4.54 KB

What I’m doing next is just preparing dummy train and evaluation, please forgive me if this code hurts your eyes 
    model = UNet(nb_classes=6)
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    #summary(model, input_size=(3, 128, 128))


    for t in range(99):
        image = load_image(f'data/images/image_{t}.jpg')
        mask = np.load(f'data/masks/mask_{t}.npy')

        # loading images as 3 channel RGB
        image = np.expand_dims(image, 0)
        image = image.transpose(0, 3, 1, 2)
        input_data = torch.tensor(image)

        # loading masks as numpy 6 channel arrays
        mask_data = torch.tensor(mask)

        # argmax to combine all channels?
        y_true = torch.argmax(mask_data, dim=0)
        # adding one BATCH dimension
        y_true = y_true.unsqueeze(0)
        y_pred = model(input_data)

        loss = criterion(y_pred, y_true)
        print(t, loss.item())

        # Zero gradients, perform a backward pass, and update the weights.
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    model.eval()
    image = load_image('data/images/image_102.jpg')
    image = np.expand_dims(image, 0)
    image = image.transpose(0, 3, 1, 2)
    input_data = torch.tensor(image)

    pred = model(input_data)
    pred = F.sigmoid(pred)
    print(pred.shape)
    np.save('pred', pred)

Summarizing, I have following questions:

As I understand, in the multi-label segmentations layers can overlay (each pixel can be assigned to multiple objects), where in the multi-class segmentation, each pixel can be assigned only to one object?

So for the multi-class, we can have one channel mask (3x3) which can look like this:
[ 0 1 2 | 1 2 3 | 4 5 5 ]  
We would prepare targets as [ batch_size, channel=1, height, width ], please correct me if I am wrong.
Which loss function should I use in this case?


How to proceed with the multi-label segmentation?


Right now the loss after ~ 40 loops drops to around 0.05. I’m using code from the linked github:


loaded_array = np.load('pred.npy')
loaded_array.shape
pred_rgb = [helper.masks_to_colorimg(x) for x in loaded_array]
plt.imshow(pred_rgb[0])

Masks for the 102 image look like this (after 99 loops):
99 loops
What am I doing wrong?