Edit: as is turns out, not even the model's initial creator could successfully fine-tune it. This is most likely a problem of implementation, or possibly related to the non-intuitive way in which the Keras batch normalization layer works.
I'm trying to fine-tune this Keras implementation of Google's DeepLab v3+ model on a custom dataset that is derived from the non-augmented Pascal VOC 2012 benchmark dataset (1449 training examples) for my research concerns.
I figured I'd first try to just retrain it myself on the original Pascal VOC dataset and try to get results close to those of the paper. The repo's author apparently succeeded in doing this so there's little chance the Keras model is incorrect. I successfully loaded the model pre-trained on ImageNet (from Google's official model zoo) and the feature maps clearly show the model is able to distinguish all the objects in pictures it's fed with (see figures below).
I'm freezing the 356 first layers which correspond to the initial backbone (Xception in my case.) I added an additional final softmax layer in the model since the model from the article initially outputs logits. Related to this choice, the dataset does have a background class. I'm using the tf.keras.optimizers.Adadelta optimizer.
However, after weeks of tweaking and exploring, I still can't get the model to learn or do anything worthwhile when it comes to segmentation.
I've tried using dozens of different loss and accuracy functions found here and there on the web, mainly variants of pixel-wise cross-entropy and soft dice loss, as well as tweaking the learning rate from $10^{-1}$ to $10^{-5}$ (the authors used $10^{-2}$ in the original paper), and every time I get the same; the loss value basically oscillates around a fairly small value and then the early stopping callback I'm using stops the training phase after $7$ or $8$ epochs.
Here's the evolution of a typical metric if I don't stop the process (in this case the learning rate was set to $10^{-5}$, batch size of $10$):
I decided to make a prediction on the same image after each epoch and here's what it looks like after the first epoch (bottom right "labels" a picture is just the argmax on the feature maps):
And after 20 epochs:
All intermediary results look something like, apparently independently from the hyperparameters.
I even tried using the accuracy and loss functions that the repo's author said he used to do exactly what I'm trying to do, but I get the same chaotic metric curves.
I'm seriously running out of ideas on where this could be coming from. I'd love to get hints on where I should look next for a possible mistake that I made.
Dataflow details
I'm using TensorFlow's dataset API (essentially following this very good guide) to load the dataset in memory. Said dataset was beforehand shuffled and split into $140$ shards of $10$ examples, which is the maximum batch size I can use on my hardware. I then select a shuffled set of shards and preprocess the examples in them by rescaling/padding/cropping them to size $512 \times 512$ with intensity values between $-1$ and $1$, cast them to tf.float32 tensors and generate $21$ binary masks for each class of the dataset.
- an input tensor is a batch of shape $(10, 512, 512, 3)$ with values in $[-1, 1$] and encoded in
float32; - associated ground truth is a tensor of shape $(10, 512, 512, 21)$ with values being either $0$, $1$ or $255$ (the latter value is used for "ambiguous" or padded areas; in turn parts of the image to ignore).
Loss and accuracy functions
I start by ignoring the labels and predictions in ignored areas (see value $255$ above):
def get_valid_labels_and_logits(y_true, y_pred):
valid_labels_mask = tf.not_equal(y_true, 255.0)
indices_to_keep = tf.where(valid_labels_mask)
valid_labels = tf.gather_nd(params=y_true, indices=indices_to_keep)
valid_logits = tf.gather_nd(params=y_pred, indices=indices_to_keep)
return valid_labels, valid_logits
I triple checked this on a tiny custom $2 \times 3$ image and it works as expected.
Next I compute a dice loss averaged on all classes as defined in this article:
def soft_dice_loss(y_true, y_pred):
y_true, y_pred = get_valid_labels_and_logits(y_true, y_pred)
# Next tensors are of shape (num_batches, num_classes)
interception_volume = tf.reduce_sum(tf.reduce_sum(y_true * y_pred, axis=1), axis=1)
y_true_sum_per_class = tf.reduce_sum(tf.reduce_sum(y_true, axis=1), axis=1)
y_pred_sum_per_class = tf.reduce_sum(tf.reduce_sum(y_pred, axis=1), axis=1)
return tf.reduce_mean(1.0 - 2.0 * interception_volume / (y_true_sum_per_class + y_pred_sum_per_class))
I tried different variants of these, including the native cross-entropy and binary cross-entropy but it didn't change the behavior much.
The default accuracy function doesn't seem to work, so I implemented a custom mean IoU accuracy function which works well on examples I gave it by hand.
Apologies for this wall of text but I wanted to make the situation clear. Thank you so much for your kind help and advice!
