Modifying a strategy to allow for the model to have more than one output.

[JothamWong]

I am trying to implement a strategy called Domain Adversarial Training where my model needs to output both a class prediction as well as a domain prediction.
My first thought was to override the training_epoch method but I wanted to see if there was already an established best practise for this.

[vlomonaco]

Answer from @Mattdl: "self.logits is returned the output of the self.model(), hence calling the forward() if your model. By overriding the forward, you can just return a tuple in forward(), then you can access this anywhere in your strategy with self.logits[index-in-tuple]."

[JothamWong]

Yep it would work if my model's forward method signature is the same.

However in the slack message, I forgot to mention that my model's forward method would also take in an additional parameter that controlled how important the domain prediction (the second output) is. My mistake for not including this piece of information in my original post.

As such, I decided to just override the training_epoch method and replace the line
self.logits = self.model(self.mb_x)
with
self.logits, self.domain_prediction = self.model(self.mb_x, self.grl_lambda)

[vlomonaco]

This is an interesting use-case! Thank you @psychicmario for bringing it to our attention. We will discuss with @AntonioCarta (lead of the Training module) if it makes sense to add a method to the callbacks system to customize the forward pass. In the meanwhile your solution of overriding the training_epoch seems appropriate!

Please mark this answer as approved if there are not more questions on your side :)