<p>I think the problem comes from your classes. What are their formats ? Are they one-hot-encoded ?</p>
<p>Your activation and loss functions are correct for me. I think your last convolutionnal layer should be with a filter of size 1x1 (similarly to the original <a href="https://arxiv.org/abs/1505.04597" rel="nofollow noreferrer">paper</a>). It ensures better pixel-classificiation.</p>
Implementing UNet model for Multi Class (4-classes) Image Segmentation
<p>I am trying to implement a UNet model, on labeled image data. The dataset contains around 10,000 images and their respective masks (colored-RGB).</p>
<p>Image Dimensions: 500 X 500 X 3</p>
<p>The masks are not black &amp; white, they are colored (RGB), having 3 classes (technically 4):</p>
<ul>
<li>Background: Black</li>
<li>Class 1: Red</li>
<li>Class 2: Green</li>
<li>Class 3: Blue</li>
</ul>
<p>This is the code for the last two CONV blocks of the model:</p>
<pre><code>    model = Conv2D(64,(3,3),strides=(1, 1),padding='same')(concat_5)
    model = LeakyReLU(0.1)(model)
    model = BatchNormalization()(model)

    model = Conv2D(4,(3,3),strides=(1, 1),padding='same', activation=&quot;softmax&quot;)(model)

    model = Model(base_model.input,model)
</code></pre>
<p>Model Architecture:</p>
<pre><code>Model: &quot;functional_1&quot;
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            [(None, 500, 500, 3) 0                                            
__________________________________________________________________________________________________
block1_conv1 (Conv2D)           (None, 500, 500, 64) 1792        input_1[0][0]                    
__________________________________________________________________________________________________
block1_conv2 (Conv2D)           (None, 500, 500, 64) 36928       block1_conv1[0][0]               
__________________________________________________________________________________________________
block1_pool (MaxPooling2D)      (None, 250, 250, 64) 0           block1_conv2[0][0]               
__________________________________________________________________________________________________
block2_conv1 (Conv2D)           (None, 250, 250, 128 73856       block1_pool[0][0]                
__________________________________________________________________________________________________
block2_conv2 (Conv2D)           (None, 250, 250, 128 147584      block2_conv1[0][0]               
__________________________________________________________________________________________________
block2_pool (MaxPooling2D)      (None, 125, 125, 128 0           block2_conv2[0][0]               
__________________________________________________________________________________________________
block3_conv1 (Conv2D)           (None, 125, 125, 256 295168      block2_pool[0][0]                
__________________________________________________________________________________________________
block3_conv2 (Conv2D)           (None, 125, 125, 256 590080      block3_conv1[0][0]               
__________________________________________________________________________________________________
block3_conv3 (Conv2D)           (None, 125, 125, 256 590080      block3_conv2[0][0]               
__________________________________________________________________________________________________
block3_pool (MaxPooling2D)      (None, 62, 62, 256)  0           block3_conv3[0][0]               
__________________________________________________________________________________________________
block4_conv1 (Conv2D)           (None, 62, 62, 512)  1180160     block3_pool[0][0]                
__________________________________________________________________________________________________
block4_conv2 (Conv2D)           (None, 62, 62, 512)  2359808     block4_conv1[0][0]               
__________________________________________________________________________________________________
block4_conv3 (Conv2D)           (None, 62, 62, 512)  2359808     block4_conv2[0][0]               
__________________________________________________________________________________________________
block4_pool (MaxPooling2D)      (None, 31, 31, 512)  0           block4_conv3[0][0]               
__________________________________________________________________________________________________
block5_conv1 (Conv2D)           (None, 31, 31, 512)  2359808     block4_pool[0][0]                
__________________________________________________________________________________________________
block5_conv2 (Conv2D)           (None, 31, 31, 512)  2359808     block5_conv1[0][0]               
__________________________________________________________________________________________________
block5_conv3 (Conv2D)           (None, 31, 31, 512)  2359808     block5_conv2[0][0]               
__________________________________________________________________________________________________
block5_pool (MaxPooling2D)      (None, 15, 15, 512)  0           block5_conv3[0][0]               
__________________________________________________________________________________________________
conv2d_transpose_5 (Conv2DTrans (None, 31, 31, 256)  1179904     block5_pool[0][0]                
__________________________________________________________________________________________________
leaky_re_lu_10 (LeakyReLU)      (None, 31, 31, 256)  0           conv2d_transpose_5[0][0]         
__________________________________________________________________________________________________
batch_normalization_10 (BatchNo (None, 31, 31, 256)  1024        leaky_re_lu_10[0][0]             
__________________________________________________________________________________________________
concatenate_5 (Concatenate)     (None, 31, 31, 768)  0           batch_normalization_10[0][0]     
                                                                 block5_conv3[0][0]               
__________________________________________________________________________________________________
conv2d_6 (Conv2D)               (None, 31, 31, 512)  3539456     concatenate_5[0][0]              
__________________________________________________________________________________________________
leaky_re_lu_11 (LeakyReLU)      (None, 31, 31, 512)  0           conv2d_6[0][0]                   
__________________________________________________________________________________________________
batch_normalization_11 (BatchNo (None, 31, 31, 512)  2048        leaky_re_lu_11[0][0]             
__________________________________________________________________________________________________
conv2d_transpose_6 (Conv2DTrans (None, 62, 62, 512)  2359808     batch_normalization_11[0][0]     
__________________________________________________________________________________________________
leaky_re_lu_12 (LeakyReLU)      (None, 62, 62, 512)  0           conv2d_transpose_6[0][0]         
__________________________________________________________________________________________________
batch_normalization_12 (BatchNo (None, 62, 62, 512)  2048        leaky_re_lu_12[0][0]             
__________________________________________________________________________________________________
concatenate_6 (Concatenate)     (None, 62, 62, 1024) 0           batch_normalization_12[0][0]     
                                                                 block4_conv3[0][0]               
__________________________________________________________________________________________________
conv2d_7 (Conv2D)               (None, 62, 62, 512)  4719104     concatenate_6[0][0]              
__________________________________________________________________________________________________
leaky_re_lu_13 (LeakyReLU)      (None, 62, 62, 512)  0           conv2d_7[0][0]                   
__________________________________________________________________________________________________
batch_normalization_13 (BatchNo (None, 62, 62, 512)  2048        leaky_re_lu_13[0][0]             
__________________________________________________________________________________________________
conv2d_transpose_7 (Conv2DTrans (None, 125, 125, 512 2359808     batch_normalization_13[0][0]     
__________________________________________________________________________________________________
leaky_re_lu_14 (LeakyReLU)      (None, 125, 125, 512 0           conv2d_transpose_7[0][0]         
__________________________________________________________________________________________________
batch_normalization_14 (BatchNo (None, 125, 125, 512 2048        leaky_re_lu_14[0][0]             
__________________________________________________________________________________________________
concatenate_7 (Concatenate)     (None, 125, 125, 768 0           batch_normalization_14[0][0]     
                                                                 block3_conv3[0][0]               
__________________________________________________________________________________________________
conv2d_8 (Conv2D)               (None, 125, 125, 256 1769728     concatenate_7[0][0]              
__________________________________________________________________________________________________
leaky_re_lu_15 (LeakyReLU)      (None, 125, 125, 256 0           conv2d_8[0][0]                   
__________________________________________________________________________________________________
batch_normalization_15 (BatchNo (None, 125, 125, 256 1024        leaky_re_lu_15[0][0]             
__________________________________________________________________________________________________
conv2d_transpose_8 (Conv2DTrans (None, 250, 250, 256 590080      batch_normalization_15[0][0]     
__________________________________________________________________________________________________
leaky_re_lu_16 (LeakyReLU)      (None, 250, 250, 256 0           conv2d_transpose_8[0][0]         
__________________________________________________________________________________________________
batch_normalization_16 (BatchNo (None, 250, 250, 256 1024        leaky_re_lu_16[0][0]             
__________________________________________________________________________________________________
concatenate_8 (Concatenate)     (None, 250, 250, 384 0           batch_normalization_16[0][0]     
                                                                 block2_conv2[0][0]               
__________________________________________________________________________________________________
conv2d_9 (Conv2D)               (None, 250, 250, 128 442496      concatenate_8[0][0]              
__________________________________________________________________________________________________
leaky_re_lu_17 (LeakyReLU)      (None, 250, 250, 128 0           conv2d_9[0][0]                   
__________________________________________________________________________________________________
batch_normalization_17 (BatchNo (None, 250, 250, 128 512         leaky_re_lu_17[0][0]             
__________________________________________________________________________________________________
conv2d_transpose_9 (Conv2DTrans (None, 500, 500, 128 147584      batch_normalization_17[0][0]     
__________________________________________________________________________________________________
leaky_re_lu_18 (LeakyReLU)      (None, 500, 500, 128 0           conv2d_transpose_9[0][0]         
__________________________________________________________________________________________________
batch_normalization_18 (BatchNo (None, 500, 500, 128 512         leaky_re_lu_18[0][0]             
__________________________________________________________________________________________________
concatenate_9 (Concatenate)     (None, 500, 500, 192 0           batch_normalization_18[0][0]     
                                                                 block1_conv2[0][0]               
__________________________________________________________________________________________________
conv2d_10 (Conv2D)              (None, 500, 500, 64) 110656      concatenate_9[0][0]              
__________________________________________________________________________________________________
leaky_re_lu_19 (LeakyReLU)      (None, 500, 500, 64) 0           conv2d_10[0][0]                  
__________________________________________________________________________________________________
batch_normalization_19 (BatchNo (None, 500, 500, 64) 256         leaky_re_lu_19[0][0]             
__________________________________________________________________________________________________
conv2d_11 (Conv2D)              (None, 500, 500, 4)  2308        batch_normalization_19[0][0]     
==================================================================================================
Total params: 31,948,164
Trainable params: 31,941,892
Non-trainable params: 6,272
__________________________________________________________________________________________________
</code></pre>
<p>When I am training the model, I am getting the following error:</p>
<pre><code>InvalidArgumentError:  Incompatible shapes: [4,4,500,500] vs. [4,3,500,500]
     [[node categorical_crossentropy/mul (defined at &lt;ipython-input-15-b60b11ed9e76&gt;:9) ]] [Op:__inference_train_function_5874]

Function call stack:
train_function
</code></pre>
<p>It seems that there is some problem with the dimensions of the last CONV layer of the model.
I want suggestions about the following three things:</p>
<ol>
<li>Dimensions of the last CONV layer (considering multi-class segmentation for <strong>4 classes</strong>)</li>
<li>Which activation function should I use in the last CONV layer? Currently, I am using <code>softmax</code> activation.</li>
<li>What kind of loss should I use while compiling the model? Currently, I am using <code>categorical_crossentropy</code>.</li>
</ol>