Keras-yolo3: ValueError: Unknown entry in loss dictionary: "yolo_loss". Only expected the following keys: ['conv2d_209', 'conv2d_217', 'conv2d_225']

[sirius0503]

I am trying to train this model on my custom dataset( which has only a single class for 'person') using Google Colab. After the first 50 epochs of using full Yolo with the output layers frozen, freeze = 2I get GPU memory RunOutOfMemory error with Tesla T4 GPU. So, I tried to train the model till the first 50 epochs, then saved the model weights in logs/000/train_weights_stage_1.h5 and then re-load the model with the new weights as follows:

from keras.callbacks import ModelCheckpoint
from keras.optimizers import Adam
import numpy as np
from keras.layers import Input, Lambda




from yolo3.model import preprocess_true_boxes, yolo_body, tiny_yolo_body, yolo_loss
from yolo3.utils import get_random_data


import keras.losses
keras.losses.custom_loss = yolo_loss

log_dir = 'logs/000/'
annotation_path = 'model_data/annotations.txt'
val_split = 0.1

with open(annotation_path) as f:
    lines = f.readlines()
np.random.seed(10101)
np.random.shuffle(lines)
print(lines)
np.random.seed(None)
num_val = int(len(lines)*val_split)
num_train = len(lines) - num_val

checkpoint = ModelCheckpoint(log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
        monitor='val_loss', save_weights_only=True, save_best_only=True, period=10, verbose=0)



classes_path = 'model_data/new_classes.txt'


def get_classes(classes_path):
    '''loads the classes'''
    with open(classes_path) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names

class_names = get_classes(classes_path)  
num_classes = len(class_names)


from yolo3.model import yolo_head, yolo_body
from keras.layers import Input, Lambda
from keras.utils import CustomObjectScope


model = yolo_body(Input(shape=(None, None, 3)), 3, num_classes)
model.load_weights('logs/000/trained_weights_stage_1.h5', by_name=True) 


# Unfreeze and continue training, to fine-tune.
# Train longer if the result is not good.
if True:
    for i in range(len(model.layers)):
        model.layers[i].trainable = True
    model.compile(optimizer=Adam(lr=1e-4), loss={'yolo_loss': lambda y_true, y_pred: y_pred}) # recompile to apply the change
    print('Unfreeze all of the layers.')

    batch_size = 32 # note that more GPU memory is required after unfreezing the body
    print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
    #import pdb; pdb.set_trace()
    history = model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),
                  steps_per_epoch=max(1, num_train//batch_size),
                  validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),
                  validation_steps=max(1, num_val//batch_size),
                  epochs=100,
                  initial_epoch=50,
                  callbacks=[logging, reduce_lr, early_stopping]) #removed checkpointing.
    model.save_weights(log_dir + 'trained_weights_final.h5')

  # Further training if needed.`

And, receive the following error:

ValueError                                Traceback (most recent call last)
<ipython-input-30-d6473c9875b2> in <module>()
     61     for i in range(len(model.layers)):
     62         model.layers[i].trainable = True
---> 63     model.compile(optimizer=Adam(lr=1e-4), loss={'yolo_loss': lambda y_true, y_pred: y_pred}) # recompile to apply the change
     64     print('Unfreeze all of the layers.')
     65 

/usr/local/lib/python3.6/dist-packages/keras/engine/training.py in compile(self, optimizer, loss, metrics, loss_weights, sample_weight_mode, weighted_metrics, target_tensors, **kwargs)
    117                                      'dictionary: "' + name + '". '
    118                                      'Only expected the following keys: ' +
--> 119                                      str(self.output_names))
    120             loss_functions = []
    121             for name in self.output_names:

ValueError: Unknown entry in loss dictionary: "yolo_loss". Only expected the following keys: ['conv2d_209', 'conv2d_217', 'conv2d_225']

At a loss as to how to get rid of this error and train the full model with the saved weights from the first part. Also, I tried loading the full model from model.save('trained_model_stage_1.h5'), and using model.load_model('trained_model_stage_1.h5', custom_objects={'yolo_loss' : yolo_loss}, but that also doesn't work for me, and I get error of yolo_head is not defined

Would be great if you could help @qqwweee

[lapsya]

As far as I understand, the problem is that you create your model as just
model = yolo_body(Input(shape=(None, None, 3)), 3, num_classes), which has 3 outputs: ['conv2d_209', 'conv2d_217', 'conv2d_225'] , but then compile the model with loss defined for output layer called 'yolo_loss'.
You can see how it is done in the original 'train.py': a layer called 'yolo_loss' is added to the model, that gets as an input model's outputs and ground truth, and returns the loss value. This is why the model is compiled with loss loss={'yolo_loss': lambda y_true, y_pred: y_pred}
To be honest, I have no idea why it's implemented this way.

[sirius0503]

@lapsya : Can you explain the lambda usage for the layer yolo_loss , I understand from this line, that model_loss layer is named as yolo_loss and the loss is computed for this Lambda layer, but I don't understand the usage of lambda y_true, y_pred: y_pred? It seems to me that using y_true and y_pred as arguments the output is y_pred, which is the computed loss?

Thanks for your previous answer!

[lapsya]

This Lambda layer essentially computes the loss from the model's previous output (model_body.output) and additional ground truth input (y_true) and returns it as the final model output.
When you're compiling the model, you need to specify loss as either a function or a dictionary, like loss={'output_name': loss_for_this_output}. The loss function is expected to take 2 arguments: y_true and y_pred. Here the loss is passed as {'yolo_loss': lambda y_true, y_pred: y_pred}, meaning that the loss for output with the name yolo_loss is the value of the output itself. You can read more about it in the official keras docs

[sirius0503]

But, for the lambda layer named by yolo_loss, the inputs are [*model_body.output, *y_true], but the lambda is providing y_true, y_pred as the input , even if y_pred be equal to model_body.output, the position of y_true, and y_pred seem to be reversed. Can you explain why?

[lapsya]

Lambda layer named 'yolo_loss' is exactly that - just a layer of the model that happens to take inputs [*model_body.output, *y_true] and computes value that we know to be loss value.
loss={'yolo_loss': lambda y_true, y_pred: y_pred} on the other hand specifies model's loss function - a lambda that has standard loss function signature loss(y_true, y_pred).

Now the model takes as inputs [model_body.input, *y_true] and produces output model_loss, which is the output of Lambda layer 'yolo_loss'.
The model is then compiled with loss={'yolo_loss': lambda y_true, y_pred: y_pred}, which means that the model's loss function takes its output y_pred and interprets it as loss value regardless of what was passed as y_true