Move backwards step into model
jensheit opened this issue · 1 comments
In case of multiple chained models for example source separation + speech recognition it might be necessary to do intermediate backwards steps to reduce the required gpu memory during training.
The user could be enabled to use mulitple backwards steps by moving the backwards step and the train_step into the model.
However, we have to consider the implications for the Hook post_step, which is at the moment called after train_step but before the backwards step.
Another open question is how to handle the timer information.
However, we have to consider the implications for the Hook post_step, which is at the moment called after train_step but before the backwards step.
This is currently done to decrease the memory consumption (i.e. after the post step we can delete the input, and the review).
Another point to consider is, that the multi gpu source code must be changed.
I don't know if calling backward in a thread is allowed and recommended in pytorch.
I would say, we plan to implement it, when we see a demand for it.
A possible workaround (for those that need it now):
- Use single GPU
- Just call backward inside the model
- Return the last loss and the trainer calls backward on this loss
The code wouldn't be pretty, but it should do the task.