python - How to use adaptive loss function from google-research in Keras?

Question

Similar to this question, I am having some trouble using the adaptive loss function from robust loss (commit c24a2d8) from google-research: the AdaptiveLossFunction class works as a loss function, but is not 'adaptive' (latent parameters not updating).

The example given by the author was in tf1 fashion, and I cannot figure out how to make a loss class that I can compile with a Keras model and monitor alpha & scale in metric.

In adaptive.py, a class AdaptiveLossFunction is provided. But as it does not work if directly used as loss, I made a naive wrapper:

class adaLoss(AdaptiveLossFunction):

  def __init__(self, num_channels, float_dtype, name='adaLoss'):
    super().__init__(num_channels, float_dtype)
    self.__name__=name

  def __call__(self,t,p):
    return super().__call__(t-p)

  def alpha(self,t=0,p=0):
    return super().alpha()

  def scale(self,t=0,p=0):
    return super().scale()

And here's the result with some NN has only one output:

opt=Adam(lr=lr_max)
loss=adaLoss(1,np.float32)
model.compile(loss=loss, optimizer=opt, metrics=[loss.alpha,loss.scale])
model.fit(x=X_train, y=Y_train, verbose=2, validation_data=(X_val,Y_val), epochs=epochs, batch_size=batch_size, callbacks=callbacks)

Epoch 1/1454
140/140 - 4s - loss: 55.1648 - alpha: 1.0000 - scale: 1.0000 - val_loss: 12.5521 - val_alpha: 1.0000 - val_scale: 1.0000
Epoch 2/1454
140/140 - 2s - loss: 17.4579 - alpha: 1.0000 - scale: 1.0000 - val_loss: 12.1275 - val_alpha: 1.0000 - val_scale: 1.0000
Epoch 3/1454
140/140 - 2s - loss: 16.3631 - alpha: 1.0000 - scale: 1.0000 - val_loss: 10.9933 - val_alpha: 1.0000 - val_scale: 1.0000
Epoch 4/1454
140/140 - 2s - loss: 15.4100 - alpha: 1.0000 - scale: 1.0000 - val_loss: 9.2542 - val_alpha: 1.0000 - val_scale: 1.0000
Epoch 5/1454
140/140 - 2s - loss: 14.3685 - alpha: 1.0000 - scale: 1.0000 - val_loss: 7.9095 - val_alpha: 1.0000 - val_scale: 1.0000

The model is indeed training with this loss function, but alpha and scale parameters are not updated. Checking withing calling loss.trainable_variables gives me, which confirms the latent variables are not updating.

(<tf.Variable 'LatentAlpha:0' shape=(1, 1) dtype=float32, numpy=array([[-0.]], dtype=float32)>,
 <tf.Variable 'LatentScale:0' shape=(1, 1) dtype=float32, numpy=array([[0.]], dtype=float32)>)

Any help is appreciated. Thanks.

question from:https://stackoverflow.com/questions/66066179/how-to-use-adaptive-loss-function-from-google-research-in-keras

Answer 1

The tf.Variable of the loss function are not optimized by the call to fit because they don't belong to the training_variables collection of the model.

A quick and dirty way is to add the latent alpha and latent scale tf.Variable to the keras model by using the hidden property tf._trainable_weights.

Something akin to this:

opt=tf.keras.optimizers.Adam(lr=1e-3)
loss = adaLoss(1, tf.float32)
X_train = tf.random.normal((1000,1))
Y_train = X_train*3 + 4
model = tf.keras.Sequential([tf.keras.layers.Dense(1,input_shape=(1,))])

# adding the variable from the loss function to the model variables
model._trainable_weights += [loss._latent_alpha, loss._latent_scale]
model.compile(loss=loss, optimizer=opt, metrics=[loss.alpha,loss.scale])

Then, running a few epochs:

>>> model.fit(x=X_train, y=Y_train, verbose=2, epochs=3)
Epoch 1/3
32/32 - 2s - loss: 4.6627 - alpha: 0.9919 - scale: 1.0104
Epoch 2/3
32/32 - 0s - loss: 4.5371 - alpha: 0.9765 - scale: 1.0305
Epoch 3/3
32/32 - 0s - loss: 4.4217 - alpha: 0.9620 - scale: 1.0497

A more explicit way of doing this would be to write a custom training loop, and adding the two variables of the loss to the gradients calculation and to the optimizer. (See the guide: Writing a training loop from scratch)