So I'm implementing Center loss: https://ydwen.github.io/papers/WenECCV16.pdf and I am having problem with updating weights in my layer, which here means updating centers in Center loss. When I print my class_centers like this tf.print(self.class_centers, summarize=-1, output_stream='file:///tensors.txt') than they never change. When I print other Variables they seem fine, so the only problem I can think of is that add_update() doesn't do what it should do.
The custom layer:
class CenterLossLayer(Layer):
def __init__(self, alpha=0.5, **kwargs):
self.alpha = alpha
super(CenterLossLayer, self).__init__(**kwargs)
def build(self, input_shape):
print('Center loss input 1 (feature_size): ', input_shape[0][1])
print('Center loss input 2 (num_classes): ', input_shape[1][1])
self.class_centers = self.add_weight(name='class_centers',
shape=(input_shape[1][1], input_shape[0][1]),
initializer='uniform',
trainable=False)
super(CenterLossLayer, self).build(input_shape)
def call(self, x, mask=None):
embeddings, one_hots = x
tf.print(self.class_centers, summarize=-1, output_stream='file:///tensors.txt')
batch_centers = K.dot(one_hots, self.class_centers)
batch_delta = batch_centers - embeddings
class_delta = K.dot(K.transpose(one_hots), batch_delta)
counts = K.sum(K.transpose(one_hots), axis=1, keepdims=True) + 1
class_delta = class_delta / counts
class_delta = K.in_train_phase(self.alpha * class_delta, 0 * class_delta)
updated_class_centers = self.class_centers - class_delta
self.add_update((self.class_centers, updated_class_centers), x[0])
losses = K.sum(K.square(embeddings - batch_centers), axis=1, keepdims=True)
return losses
def compute_output_shape(self, input_shape):
return (input_shape[1][0], )
and the final loss is:
def batch_mean_loss(y_true, y_pred):
return K.mean(y_pred, axis=0)
where y_pred is losses from CenterLossLayer.
The weird thing is that even thought the centers are not updating, the center loss is going down with each epoch and the final model is better that the one trained only with Softmax loss.