Documentation for add_loss method of tf.keras.model

Viewed 1783

Whenever I've worked with tensorflow's keras api in the past, I've specified the loss function for a model with model.compile. I'm currently working on a repo which uses 'add_loss' to specify the loss function inside of model.call. Or at least, that's what I am assuming is happening, because I can't find any documentation for this method (i.e. none on https://www.tensorflow.org/api_docs/python/tf/keras/Model), and I can't find any tutorials that use this method either. What's more, I can't even figure out where it's defined in the source code.

class TRPO(Model):
    def __init__(self, obs_dim, act_dim, hid1_mult, kl_targ, init_logvar, eta, **kwargs):
        super(TRPO, self).__init__(**kwargs)
        self.kl_targ = kl_targ
        self.eta = eta
        self.beta = self.add_weight('beta', initializer='zeros', trainable=False)
        self.policy = PolicyNN(obs_dim, act_dim, hid1_mult, init_logvar)
        self.logprob = LogProb()
        self.kl_entropy = KLEntropy()

    def call(self, inputs):
        obs, act, adv, old_means, old_logvars, old_logp = inputs
        new_means, new_logvars = self.policy(obs)
        new_logp = self.logprob([act, new_means, new_logvars])
        kl, entropy = self.kl_entropy([old_means, old_logvars,
                                       new_means, new_logvars])
        loss1 = -K.mean(adv * K.exp(new_logp - old_logp))
        loss2 = K.mean(self.beta * kl)
        # TODO - Take mean before or after hinge loss?
        loss3 = self.eta * K.square(K.maximum(0.0, K.mean(kl) - 2.0 * self.kl_targ))
        self.add_loss(loss1 + loss2 + loss3)

        return [kl, entropy]

Anyone have experience with using add_loss and can point out how it works? And explain why you wouldn't just write your own loss function and pass that in model.compile?

1 Answers

You can find the official documentation of add_loss here. Add loss tensor(s), potentially dependent on layer inputs. This method can be used inside a subclassed layer or model's call function, in which case losses should be a Tensor or list of Tensors. There are few example in the documentation to explain the add_loss.

You can find the source code of add_loss in tf.keras.layers.Layer. This is the class from which all layers inherit. Click on "View source on GitHub" and search for add_loss.

Coming to your questions -

Anyone have experience with using add_loss and can point out how it works?

You can find good example using add_loss here and here with explanations.

And explain why you wouldn't just write your own loss function and pass that in model.compile?

model.compile() loss functions in Tensorflow always take two parameters y_true and y_pred. Using model.add_loss() has no such restriction and allows you to write much more complex losses that depend on many other tensors, but it has the inconvenience of being more dependent on the model, whereas the standard loss functions work with just any model.

Hope this answers your question. Happy Learning.

Related