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?