[Bug] loss_fn argument for Trainer must not be a function since 1.2.4 about trax HOT 6 CLOSED

I think the way we're trying to go is trax.layers.Fn('MyLoss', loss_fn) - Fn is just Lambda. Thanks for your help as we try to make it understandable - we need to improve the docs too!

Closing for now as the immediate issue is resolved.

from trax.

Let me look into this, but if possible can you send a repro? Is the problem in 1.2.4 or earlier?

from trax.

This was changed on purpose in 1.2.4 so as to make sure losses are indeed just layers. A quick fix here is just making it CrossEntropyLoss().

But let's think if the quick fix is enough. Stefan: do you have a use-case where having functions-creating-layers instead of layer instances is actually desireable?

from trax.

Hi! :)

I only intended to point this out. E.g. is the naming loss_fn a bit misleading. Also the old examples are passing only the type loss_fn=trax.layers.CrossEntropyLoss instead of the object.

@afrozenator The problem starts with 1.2.4 and it's easy to reproduce by just passing something that is not a layer to loss_fn

from trax.

Thanks for pointing this out! I agree the naming is now misleading - and sorry for the incompatibility! But if there's no reason to have function there, we should probably go with just having layers and maybe just calling it loss? Thanks for bringing this to our attention!

from trax.

You're welcome. No big deal - in cases like this exceptions are rather clear and one can figure out what to do :) Regarding the argument: I really know too little at this point about trax to be able to say what makes sense but one idea would be to do like Keras when compiling the model (see docs):

Arguments:

• loss: String (name of objective function), objective function or tf.keras.losses.Loss instance. See tf.keras.losses. An objective function is any callable with the signature scalar_loss = fn(y_true, y_pred). If the model has multiple outputs, you can use a different loss on each output by passing a dictionary or a list of losses. The loss value that will be minimized by the model will then be the sum of all individual losses.

But of course that will make the implementation of Trainer a bit more complex.

One advantage of functions could be that I can just provide a function in cases where I already have one in my code. Also I don't need to write a wrapper-object for such a function or any function which does not require data or a state during execution.

After all I would say this is a design question. We can always have a trax.layers.LambdaLoss as in LambdaLoss(loss_fn=my_loss_fn) to provide a simple wrapper.

Thanks for bringing this to our attention!

I say thanks for sharing t2t and trax with us :)

from trax.