title: How to implement multi-class semantic segmentation?
body: <p>I'm able to train a <a href="https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/" rel="noreferrer">U-net</a> with labeled images that have a binary classification. </p>
  
  <p>But I'm having a hard time figuring out how to configure the final layers in Keras/Theano for <strong>multi-class</strong> classification (4 classes).</p>
  
  <p>I have 634 images and corresponding 634 masks that are <code>unit8</code> and 64 x 64 pixels.</p>
  
  <p>My masks, instead of being black (0) and white (1), have color labeled objects in 3 categories plus background as follows:</p>
  
  <ul>
  <li>black (0), background</li>
  <li>red (1), object class 1</li>
  <li>green (2), object class 2</li>
  <li>yellow (3), object class 3</li>
  </ul>
  
  <p>Before training runs, the array containing masks is one-hot encoded as follows:</p>
  
  <pre><code>mask_train = to_categorical(mask_train, 4)
  </code></pre>
  
  <p>This makes <code>mask_train.shape</code> go from <code>(634, 1, 64, 64)</code> to <code>(2596864, 4)</code>.</p>
  
  <p>My model closely follows the Unet architecture, however the final layers seem problematic, as I'm unable to flatten the structure so as to match the one-hot encoded array.</p>
  
  <pre><code>[...]
  up3 = concatenate([UpSampling2D(size=(2, 2))(conv7), conv2], axis=1)
  conv8 = Conv2D(128, (3, 3), activation='relu', padding='same')(up3)
  conv8 = Conv2D(128, (3, 3), activation='relu', padding='same')(conv8)
  
  up4 = concatenate([UpSampling2D(size=(2, 2))(conv8), conv1], axis=1)
  conv9 = Conv2D(64, (3, 3), activation='relu', padding='same')(up4)
  conv10 = Conv2D(64, (3, 3), activation='relu', padding='same')(conv9)
  
  # here I used number classes = number of filters and softmax although
  # not sure if a dense layer should be here instead
  conv11 = Conv2D(4, (1, 1), activation='softmax')(conv10)
  
  model = Model(inputs=[inputs], outputs=[conv11])
  
  # here categorical cross entropy is being used but may not be correct
  model.compile(optimizer='sgd', loss='categorical_crossentropy',
                metrics=['accuracy'])
  
  return model
  </code></pre>
  
  <p>Do you have any suggestions on how to modify the final portions of the model so this trains successfully? I get a variety of shape mismatch errors, and the few times I managed to make it run, the loss did not change throughout epochs.</p>