Customized forward function in the base strategy

[HamedHemati]

Feature request:

The base strategy always uses avalanche_forward inside its forward function. This can be limiting when the strategy's model has more than one output, for example, (hidden_feats, prediction). Also, it feeds the task_labels to the model only if it is an instance of MultiTaskModule, so if the model is single-head, it still has to inherit from MultiTaskModule to receive task_labels (for any purpose like task-based conditioning). There could be other cases where post-processing the model's output is needed.

One solution would be defining an attribute called something like foward_func in the constructor of the base strategy initialized by its arguments and setting it as avalanche_forward by default. Then the main forward function would change as below:

    def forward(self):
        """Compute the model's output given the current mini-batch."""
        return self.forward_func(self.model, self.mb_x, self.mb_task_id)

What do you think about adding this change?

[AntonioCarta]

This is already possible, either by subclassing:

class MyStrat(Naive):
    def forward(self):
          return my_forward(...)

or, in a somewhat hackish, but perfactly working way, by overriding the forward method:

strat = Naive(...)
strat.forward = lambda self: my_foward(...)

[HamedHemati]

@AntonioCarta True, thanks.
The problem with the first solution is that you have to create a sub-class only for the forward function. So if you have 5 existing strategies, you would have to create a "MyXStrategy" version for each.

The hack looks good though but it's still a hack. I'm going with this one for now.

This is not something urgent for now but could be considered for the "general" BaseStrategy" in the future since a strategy can have different forward behaviors depending on the model.

[HamedHemati]

We can close this issue for now.

[vlomonaco]

Should I move it to the Discussions tab instead?

[AntonioCarta]

The problem with the first solution is that you have to create a sub-class only for the forward function. So if you have 5 existing strategies, you would have to create a "MyXStrategy" version for each.

IMO the solution for this problem already exists. Strategies are implemented as plugins, so you can add them to your custom loop without subclassing N times.

The problem that I see with your proposals is that we are adding a bit of indirection without much benefits. And it's not just one function. If you do it for forward then you probably want to do it for criterion, make_train_dataloader, make_eval_dataloader, and so on. If you sum all of them we have a lot of attributes and constructor arguments. I would prefer to make it more obvious which method are safe to subclass.

[HamedHemati]

Yea, this should work perfectly for the strategies implemented as a plugin, but not all strategies are/can be implemented as plugins.
I think it's only needed for the forward and criterion functions, but I'm not sure if make_train_dataloader and make_eval_dataloader need any changes as you define the output tuple in the benchmark's dataset.

In the end, this would be a question of which method would be more convenient for the user and would also allow more flexibility with less code:
(1) letting the user define and set the forward and criterion functions for existing strategies (this would need an additional callback)
or
(2) asking them to subclass as you explained (but this needs to be redone for every strategy that is not defined as a plugin)

@vlomonaco I think I should've created a discussion instead of an issue. We can have a discussion with potential features for the base strategy and keep this one on the list for later consideration if it is interesting enough.

[vlomonaco]

It is now a discussion :)