Add gradient accumulation logic to SupervisedTrainer

[jak0bw]

https://github.com/Project-MONAI/MONAI/issues/6100

Description

A few sentences describing the changes proposed in this pull request.

Types of changes

• [x] Non-breaking change (fix or new feature that would not break existing functionality).
• [ ] Breaking change (fix or new feature that would cause existing functionality to change).
• [ ] New tests added to cover the changes.
• [ ] Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• [ ] Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• [x] In-line docstrings updated.
• [ ] Documentation updated, tested make html command in the docs/ folder.

[jak0bw]

(Source code is strongly (and shamelessly) influenced by https://pytorch.org/ignite/generated/ignite.engine.supervised_training_step.html

[jak0bw]

Needs Feedback if this feature is desired and if yes probably test(s) for the new parameter.

[Nic-Ma]

Hi @jak0bw ,

Thanks for your idea and contribution here. The design of SupervisedTrainer follows the ignite engine logic, it just defines the standard (or default) computation logic for every iteration. If you have any customized logic, please pass it as a callback function: https://github.com/Project-MONAI/MONAI/blob/dev/monai/engines/trainer.py#L101

Thanks.

[jak0bw]

Hi @Nic-Ma,

thank you for your answer.

I am sorry if I misunderstand something critical here but doesn't the standard ignite logic include gradient accumulation similarly to my proposed changes (since ignite 0.4.7)? (My code is more or less copy pasted from the ignite source code) Therefore a change in the way outlined as in this pr would just restore feature parity between ignite and monai and not be considered customized logic.

Link to ignite create_supervised_trainer: https://github.com/pytorch/ignite/blob/c7c0df0fbfdff2a86415476cf0e68f36a089c1d2/ignite/engine/init.py#L404

Link to the used step function(s): https://github.com/pytorch/ignite/blob/c7c0df0fbfdff2a86415476cf0e68f36a089c1d2/ignite/engine/init.py#L44

[Nic-Ma]

Hi @jak0bw ,

Oh, I didn't notice this new option in ignite. @vfdev-5 @wyli Do you think it's necessary to add it in MONAI trainer?

Thanks in advance.

[wyli]

I think if ignite's supervised_training_step is not directly usable, we should create a util function in monai.engine.utils:

def grad_accumulation_iteration(steps=...):
    def iteration(engine, ...):
        ...
        return engine.output
    return iteration

and the usage would be

monai.engine.SupervisedTrainer(..., iteration_update=monai.engine.utils.grad_accumulation_iteration(steps), ...)

[jak0bw]

As mentioned in #6100 it is possible to directly use Ignite's supervised_training_step but as it does not emit the same Events as the monai step function some monai handlers using these Events are not triggered correctly.

[wyli]

sure, please consider creating a function in monai.engines.utils and the engine will prefer this iteration_update if it's provided: https://github.com/Project-MONAI/MONAI/blob/e375f2a17c098d7b802e5ca64322db6ce874a3aa/monai/engines/workflow.py#L125-L128

this is how we create various iteration_update functions, for example: https://github.com/Project-MONAI/MONAI/blob/dev/monai/apps/deepedit/interaction.py#LL26C7-L26C18

usage: https://github.com/Project-MONAI/tutorials/blob/aa4ca78d3e7f08c6d8f5a5a009d5da508acdb6ad/deepedit/ignite/train.py#L250C37-L259

[jak0bw]

@wyli I think I added a custom iteration update function as requested. The code still has a circular import error (the circular import of SupervisedTrainer) which I don't know how to resolve really as it kinda depends on how the monai project is structured (or deals with these problems in code).

Additionally tests are still missing.

@Nic-Ma Unfortunately, I don't really have time to further work on this pull request in the near future. Therefore this pull request (and/or the corresponding issue) can be marked for contribution wanted, closed or however the monai team wants to deal with it.