tflite_convert a Keras h5 model which has a custom loss function results in a ValueError, even if I add it in the Keras losses import

Question

I have written a SRGAN implementation. In the entry point class of the Python program, I declare a function which returns a mean square using the VGG19 model:

# <!--- COST FUNCTION --->
def build_vgg19_loss_network(ground_truth_image, predicted_image):
    loss_model = Vgg19Loss.define_loss_model(high_resolution_shape)
    return mean(square(loss_model(ground_truth_image) - loss_model(predicted_image)))

import keras.losses
keras.losses.build_vgg19_loss_network = build_vgg19_loss_network
# <!--- /COST FUNCTION --->

(Vgg19Loss class shown further below)

As you can see, I have added this custom loss function in the import keras.losses. Why? Because I thought it could solve the following problem...: When I execute the command tflite_convert --output_file=srgan.tflite --keras_model_file=srgan.h5, the Python interpreter raises this error:

raise ValueError('Unknown ' + printable_module_name + ':' + object_name) ValueError: Unknown loss function:build_vgg19_loss_network

However, it didn't solve the problem. Any other solution which could work?

Here is the Vgg19Loss class:

from keras import Model
from keras.applications import VGG19


class Vgg19Loss:
    def __init__(self):
        pass

    @staticmethod
    def define_loss_model(high_resolution_shape):
        model_vgg19 = VGG19(False, 'imagenet', input_shape=high_resolution_shape)
        model_vgg19.trainable = False
        for l in model_vgg19.layers:
            l.trainable = False
        loss_model = Model(model_vgg19.input, model_vgg19.get_layer('block5_conv4').output)
        loss_model.trainable = False
        return loss_model

Answer 1

I tried the code you posted the following way:

from keras import Model
from keras.applications import VGG19
import keras.backend as K


class Vgg19Loss:
    def __init__(self):
        pass

    @staticmethod
    def define_loss_model(high_resolution_shape):
        model_vgg19 = VGG19(False, 'imagenet', input_shape=high_resolution_shape)
        model_vgg19.trainable = False
        for l in model_vgg19.layers:
            l.trainable = False
        loss_model = Model(model_vgg19.input, model_vgg19.get_layer('block5_conv4').output)
        loss_model.trainable = False
        return loss_model


def build_vgg19_loss_network(ground_truth_image, predicted_image):
    loss_model = Vgg19Loss.define_loss_model(high_resolution_shape) # where is this variable coming from?
    return K.mean(K.square(loss_model(ground_truth_image) - loss_model(predicted_image)))

import keras.losses
keras.losses.build_vgg19_loss_network = build_vgg19_loss_network

print(keras.losses.build_vgg19_loss_network)  # <function build_vgg19_loss_network at 0x7f05e8e1cbf8>

I get no error messages and the function is assigned to the losses module. That means the problematic lines are probably not part of what you posted. It would be nice to know which line of code raises the error that you quoted.

However, I'm not sure where this high_resolution_shape argument on line 22 in your build_vgg_19_network function is coming from. If this is a global constant, it should be written in all uppercase letters separated by underscores to prevent confusion. If it is not defined it will throw a NameError sooner or later.

If I execute keras.losses.build_vgg19_loss_network(None, None) after running the code above, I get the following error message:

NameError: name 'high_resolution_shape' is not defined

Edit: If this error happens only during TFLite conversion, it does so because custom objects are not yet supported by the TFLiteConverter in tensorflow 1.x. However, there is a commit on the tensorflow Github repo that addresses this issue and adds support for custom objects (see also the related pull request). It should be part of the official tensorflow v2.0.0-beta1.