2
$\begingroup$

I'm trying to train a semantic segmentation model based on this architecture, using this one as a base. The base model uses about 10 ReLU activations, and when implemented according to the first paper, the number jumps up to 14.

The input images are of dimensions 216 x 64, and the output labels can be one of 8 classes.

Here's the complete model implementation.

I've written a custom training step, since the paper calls for it:

@tf.function
def train_step(batch_size, x_batch, y_batch, loss_func):
    with tf.GradientTape() as tape:

        # print(y_batch_train.shape)
        
        logits_strong, logits_weak = model(x_batch, training=True)  # Logits for this minibatch
        # logits = tf.concat([logits_strong[0:batch_size//2], logits_weak[batch_size//2:batch_size]], 0)
        loss_strong_value = loss_func(y_batch[0:batch_size//2], logits_strong[0:batch_size//2])
        loss_weak_value = loss_func(y_batch[batch_size//2:batch_size], logits_weak[batch_size//2:batch_size])
        loss_value = loss_strong_value + loss_weak_value
        
        # loss_value = loss_func(y_batch, logits)
        # tf.print(loss_value.shape)
        grads = tape.gradient(loss_value, model.trainable_weights)
        optimizer.apply_gradients(zip(grads, model.trainable_weights))

        train_acc_metric.update_state(y_batch[0:batch_size//2], logits_strong[0:batch_size//2])
        train_acc_metric.update_state(y_batch[batch_size//2:], logits_weak[batch_size//2:])
        # train_acc_metric.update_state(y_batch, logits)
        return loss_value

def train(model, start_epoch, num_epochs,train_dataset, optimizer, model_path, train_acc_metric, loss_fn=customized_loss, model_weights=None,):
    """
    Run a for loop with number of epochs. Run an inner for loop for each minibatch and get logits_strong and logits_weak. 
    Drop second half of logits_strong, and first half of logits_weak. Compute cross entropy loss separately and add.
    Finally, compute grads and apply. 
    
    Save model and weights after every 20 or so epochs.
    Save losses and acc for each epoch and plot after epochs are done.

    NOTE: All minibatches need to contain strong labels for first half and weak labels for second half. DO NOT SHUFFLE.
    
    Parameters: model, start_epoch, no. of epochs, optimizer, path to model, metric for train acc, model weights, loss func.

    """
    train_acc=[]
    batch_size=16
    epochs = num_epochs
    end_epoch=start_epoch + num_epochs

    if model_weights:
        load_status = model.load_weights(model_path + f"/weights/{model_weights}")
        load_status.assert_consumed()
    


    for epoch in range(start_epoch, end_epoch):
        print(f"\nStart of epoch {epoch}")
        start = time.time()
        # Iterate over the batches of the dataset.
        for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):

            loss_value = train_step(batch_size, x_batch_train, y_batch_train, loss_fn)
            

            # print(loss_value.shape, len(model.trainable_weights))

            # Log every 200 batches.
            if step % 5 == 0:
                print(
                    "Training loss (for one batch) at step %d: %.4f"
                    % (step, float(np.sum(loss_value)))
                )
                print("Seen so far: %s samples" % ((step + 1) * batch_size))
                
        
        train_acc_epoch = train_acc_metric.result()
        train_acc.append(train_acc_epoch)
        print("Training acc over epoch: %.4f" % (float(train_acc_epoch),))
        print("Time taken: %.2fs" % (time.time() - start))

        # Reset training metrics at the end of each epoch
        train_acc_metric.reset_states()
        if epoch % 10 == 0:
            model.save_weights(model_path + f"/weights/ckpt_DB_{start_epoch}_{end_epoch}")

Here are the optimizer and loss details:

optimizer = tf.keras.optimizers.SGD(learning_rate = 0.001, momentum = 0.9, nesterov = True)

#Calculation of the dice co-efficient based on actual and predicted labels
def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

def dice_coef_loss(y_true, y_pred):
    return 1-dice_coef(y_true, y_pred)

#Combined loss of weighted multi-class logistic loss and dice loss
def customized_loss(y_true,y_pred):
    # print("Shape of ground truth:", y_true.shape)
    # print("Shape of prediction:", y_pred.shape)
    return (1*K.categorical_crossentropy(y_true, y_pred))+(0.5*dice_coef_loss(y_true, y_pred)) # + 0.01*np.linalg.norm())

When I try to train it (on a relatively small dataset), this is the output:

Start of epoch 0
Training loss (for one batch) at step 0: 566345.3750
Seen so far: 16 samples
Training loss (for one batch) at step 5: 1526504.7500
Seen so far: 96 samples
Training loss (for one batch) at step 10: 1538868.5000
Seen so far: 176 samples
Training loss (for one batch) at step 15: 1445873.7500
Seen so far: 256 samples
Training loss (for one batch) at step 20: 1514306.7500
Seen so far: 336 samples
Training loss (for one batch) at step 25: 1492221.5000
Seen so far: 416 samples
Training loss (for one batch) at step 30: 1438761.3750
Seen so far: 496 samples
Training acc over epoch: 0.8664
Time taken: 13.09s

Start of epoch 1
Training loss (for one batch) at step 0: 1411657.2500
Seen so far: 16 samples
Training loss (for one batch) at step 5: 1526504.7500
Seen so far: 96 samples
Training loss (for one batch) at step 10: 1538868.5000
Seen so far: 176 samples
Training loss (for one batch) at step 15: 1445873.7500
Seen so far: 256 samples
Training loss (for one batch) at step 20: 1514306.7500
Seen so far: 336 samples
Training loss (for one batch) at step 25: 1492221.5000
Seen so far: 416 samples
Training loss (for one batch) at step 30: 1438761.3750
Seen so far: 496 samples
Training acc over epoch: 0.8944
Time taken: 10.71s

Start of epoch 2
Training loss (for one batch) at step 0: 1411657.2500
Seen so far: 16 samples
Training loss (for one batch) at step 5: 1526504.7500
Seen so far: 96 samples
Training loss (for one batch) at step 10: 1538868.5000
Seen so far: 176 samples
Training loss (for one batch) at step 15: 1445873.7500
Seen so far: 256 samples
Training loss (for one batch) at step 20: 1514306.7500
Seen so far: 336 samples
Training loss (for one batch) at step 25: 1492221.5000
Seen so far: 416 samples
Training loss (for one batch) at step 30: 1438761.3750
Seen so far: 496 samples
Training acc over epoch: 0.8944
Time taken: 10.69s

Every epoch after 2 until 20 gives the same loss and accuracy for the same batches. I'm also suspicious of the high accuracy right from the get-go. I've already tried reducing the lr to 0.001, but no change.

Could this be a divergence issue? Dying ReLU? Most importantly, how can I fix it?

$\endgroup$

1 Answer 1

4
$\begingroup$

Dying ReLU is a good guess. But the question is why this is not happening in original paper. Answer is input data features. Have you normalize the data? You can debug this issue by trying to find the original dataset from the paper and try to train your model with it.

$\endgroup$
4
  • $\begingroup$ I have normalized the data such that each pixel value is in the range [0,1]. I'm not sure if I testing the model on the data mentioned in the original paper will be possible? The base model used in the architecture is VGG16, while I'm trying it with the Dilated ReLayNet model as base. Funnily enough, the Dilated ReLayNet doesn't have this issue, even though the same dataset is being used. What could be the problem then? $\endgroup$
    – Ad Ve
    Dec 13, 2021 at 13:55
  • $\begingroup$ There are different kind of normalizations. If you are certain this is not the problem then i would suggest try to recreate the problem with a more compact code. The code you shared is hard to understand and not ready to run. $\endgroup$
    – Enes Kuz
    Dec 13, 2021 at 18:21
  • $\begingroup$ Got it! Turns out, the normalization wasn't the issue. I didn't pass the sample weights while training, which meant the class labels were very unbalanced. Passing them solved the issue. Since it was related to your answer, I'll accept it. Thanks for the help! $\endgroup$
    – Ad Ve
    Dec 13, 2021 at 19:24
  • $\begingroup$ I think this is my first accepted answer here : ) I am glad to help $\endgroup$
    – Enes Kuz
    Dec 14, 2021 at 9:07

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.