Project-MONAI / Project-MONAI/MONAI

Default for SupervisedTrainer slows down training almost 2x

Open
#8,541 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
8.7k
Forks
1.6k
Avg merge
5d 1h
Merged PRs (30d)
20

Description

https://github.com/Project-MONAI/MONAI/blob/d388d1c6fec8cb3a0eebee5b5a0b9776ca59ca83/monai/engines/trainer.py#L121-L123

The default behaviour for the SupervisedTrainer slows down each training step quite a bit. The cause of this slow-down is not apparent until profiling each step. With larger batch sizes, the delay from the default decollation step increases as more tensors have to move from GPU/CPU and back.

Image In this flame-chart, there are two steps. A forward, backward and the decollate.

It may be better to set the default for decollation in the supervised trainer to False or to add a warning that the default behaviour can significantly impact training time (in our case, 2x speed up from 4 days to 2 days by disabling decollation).

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at monai/engines/trainer.py#L121-L123 and inspect how SupervisedTrainer configures the default decollation behavior. Profile or reproduce a training step with decollation enabled and disabled, then verify that the chosen default or warning addresses the reported slowdown without changing the intended trainer behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.