2
$\begingroup$

I am currently building a U-NET semantic segmentation model on Tensorflow Keras to classify pixels as belonging to or not belonging to a class.

For this problem, I've isolated the masks for only one class (even if multiple mask classes exist) because my aim is for the model to excel at recognizing the presence of the one class in unseen images.

In the current training regimen, it appears that the model's accuracy plateaus after maybe 20 epochs of training. However, the mean Intersection over Union appears to be flat from the very first epoch to the last.

The validation accuracy ends up around 70% when the model is done training, but I suspect this is because of imbalanced classes more than that my model is any good at classifying the image pixels.

I'm relatively new to machine learning and deep learning, so really I'm hoping to get some high-level advice on how you'd typically go about solving this problem.

  • Is my problem that I've misspecified something in the UNET architecture? I chose a sigmoid activation function in the output layer and am using the Adam optimizer and binary crossentropy in the compile call. Metrics in the compile call are currently 'accuracy' and the keras.metrics.meanIoU with num_classes = 2.
  • Is Intersection Over Union even the metric I should be optimizing for in this kind of problem?
  • Are there common-sense image augmentation steps I should be attempting to improve my model's classification efficacy?
$\endgroup$

1 Answer 1

2
$\begingroup$

You could directly optimize the mean IoU loss by implementing the following loss:

def mean_iou(y_pred, y_true):
    if y_pred.shape.ndims > 1:
        y_pred = array_ops.reshape(y_pred, [-1])

    if y_true.shape.ndims > 1:
        y_true = array_ops.reshape(y_true, [-1])
    intersection = K.sum(K.sum(K.abs(y_true * y_pred), axis=-1))
    union = K.sum(K.sum(K.abs(y_true) + K.abs(y_pred), axis=-1))

    return intersection / union

Another useful loss is the Jaccard loss, which comes in handy if the segmentation classes are unbalanced:

def jaccard_distance_loss(y_true, y_pred, smooth=10):
  """
  Jaccard = (|X & Y|)/ (|X|+ |Y| - |X & Y|)
        = sum(|A*B|)/(sum(|A|)+sum(|B|)-sum(|A*B|))

  """
  intersection = K.sum(K.sum(K.abs(y_true * y_pred), axis=-1))
  union = K.sum(K.sum(K.abs(y_true) + K.abs(y_pred), axis=-1))
  jac = (intersection + smooth) / (union - intersection + smooth)
  return (1 - jac) * smooth
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.