How to implement my own loss function for Prototype learning using Keras Model

Question

I'm trying to migrate this code, "Omniglot Character Set Classification Using Prototypical Network", into Tensorflow 2.1.0 and Keras 2.3.1.

My problem is about how to use euclidean distance between train data and validation data. Look at this code:

def convolution_block(inputs, out_channels, name='conv'):

    conv = tf.layers.conv2d(inputs, out_channels, kernel_size=3, padding='SAME')
    conv = tf.contrib.layers.batch_norm(conv, updates_collections=None, decay=0.99, scale=True, center=True)
    conv = tf.nn.relu(conv)
    conv = tf.contrib.layers.max_pool2d(conv, 2)

    return conv


def get_embeddings(support_set, h_dim, z_dim, reuse=False):

    net = convolution_block(support_set, h_dim)
    net = convolution_block(net, h_dim)
    net = convolution_block(net, h_dim) 
    net = convolution_block(net, z_dim) 
    net = tf.contrib.layers.flatten(net)

    return net



support_set_embeddings = get_embeddings(tf.reshape(support_set, [num_classes * num_support_points, img_height, img_width, channels]), h_dim, z_dim)

embedding_dimension = tf.shape(support_set_embeddings)[-1]

class_prototype = tf.reduce_mean(tf.reshape(support_set_embeddings, [num_classes, num_support_points, embedding_dimension]), axis=1)

query_set_embeddings = get_embeddings(tf.reshape(query_set, [num_classes * num_query_points, img_height, img_width, channels]), h_dim, z_dim, reuse=True)


def euclidean_distance(a, b):

    N, D = tf.shape(a)[0], tf.shape(a)[1]
    M = tf.shape(b)[0]
    a = tf.tile(tf.expand_dims(a, axis=1), (1, M, 1))
    b = tf.tile(tf.expand_dims(b, axis=0), (N, 1, 1))
    return tf.reduce_mean(tf.square(a - b), axis=2)


distance = euclidean_distance(query_set_embeddings,class_prototype)

predicted_probability = tf.reshape(tf.nn.log_softmax(-distance), [num_classes, num_query_points, -1])

loss = -tf.reduce_mean(tf.reshape(tf.reduce_sum(tf.multiply(y_one_hot, predicted_probability), axis=-1), [-1]))
accuracy = tf.reduce_mean(tf.to_float(tf.equal(tf.argmax(predicted_probability, axis=-1), y)))

train = tf.train.AdamOptimizer().minimize(loss)

If I have understood everything correctly, it get the embeddings from the support_set (aka training data), and the embeddings from the query_set (aka validation data). Compute the mean for all the embeddings from the support_set, because all of them are from the same class. Then, it uses this mean to compute the distance between the embeddings from the query_set and this mean (aka class_prototype).

So, if I want to use VGG16 as the get_embeddings function. In other words, I'm going to use it to get the embeddings for the support_set and query_set:

def vgg16_feature_extractor(input_size = (200,200,1)):
    inputs = Input(input_size, name = 'input')

    conv1 = Conv2D(64, (3, 3), activation = 'relu', padding = 'same', name ='conv1_1')(inputs)
    conv1 = Conv2D(64, (3, 3), activation = 'relu', padding = 'same', name ='conv1_2')(conv1)
    pool1 = MaxPooling2D(pool_size = (2,2), strides = (2,2), name = 'pool_1')(conv1)

    conv2 = Conv2D(128, (3, 3), activation = 'relu', padding = 'same', name ='conv2_1')(pool1)
    conv2 = Conv2D(128, (3, 3), activation = 'relu', padding = 'same', name ='conv2_2')(conv2)
    pool2 = MaxPooling2D(pool_size = (2,2), strides = (2,2), name = 'pool_2')(conv2)

    conv3 = Conv2D(256, (3, 3), activation = 'relu', padding = 'same', name ='conv3_1')(pool2)
    conv3 = Conv2D(256, (3, 3), activation = 'relu', padding = 'same', name ='conv3_2')(conv3)
    conv3 = Conv2D(256, (3, 3), activation = 'relu', padding = 'same', name ='conv3_3')(conv3)
    pool3 = MaxPooling2D(pool_size = (2,2), strides = (2,2), name = 'pool_3')(conv3)

    conv4 = Conv2D(512, (3, 3), activation = 'relu', padding = 'same', name ='conv4_1')(pool3)
    conv4 = Conv2D(512, (3, 3), activation = 'relu', padding = 'same', name ='conv4_2')(conv4)
    conv4 = Conv2D(512, (3, 3), activation = 'relu', padding = 'same', name ='conv4_3')(conv4)
    pool4 = MaxPooling2D(pool_size = (2,2), strides = (2,2), name = 'pool_4')(conv4)

    conv5 = Conv2D(512, (3, 3), activation = 'relu', padding = 'same', name ='conv5_1')(pool4)
    conv5 = Conv2D(512, (3, 3), activation = 'relu', padding = 'same', name ='conv5_2')(conv5)
    conv5 = Conv2D(512, (3, 3), activation = 'relu', padding = 'same', name ='conv5_3')(conv5)
    pool5 = MaxPooling2D(pool_size = (2,2), strides = (2,2), name = 'pool_5')(conv5)

    model = Model(inputs = inputs, outputs = pool5, name = 'vgg-16_feature_extractor')

    return model

And, then on train.py:

model = vgg16_feature_extractor(input_size = (200,200,1))
model.compile(optimizer=opt, loss=my_own_loss_function, metrics=['accuracy'])
model.fit(...)

I don't know how to implement the my_own_loss_function, because this function will have only two parameters y_true, y_pred and the y_pred has to been calculated using the euclidean distance between support_set embeddings and query_set embeddings.

How do I have to implement the my_own_loss_function to use it as I want?

Maybe, y_true are the embeddings from support_set and y_pred are the embeddings from query_set.

Answer 1

Well, there are several ways you can do that.

A quite powerful solution is to define a pred layer

class PredLayer(Layer):
    """
        Layer object to calculate distance between query_embeddings and supposrt embeddings.
    """
    def __init__(self, **kwargs):
        super(PredLayer, self).__init__(**kwargs)

    def euclidean_distance(self, inputs):
        """
            Euclidean square distance.
        """
        support, query = inputs
        output = K.mean(K.square(support - query), axis=-1)
        output = K.expand_dims(output, 1)
        return output

    def call(self, inputs):
        y_pred = self.euclidean_distance(inputs)
        return y_pred

Hence you have to compose your Keras network such that your support and query embeddings will be the inputs to this layer.

...

model_pred = Model(inputs = inputs, outputs = predlayer, name = 'vgg-16_feature_extractor_pred')

return model, model_pred
```