Add AMP to ImageNet classification and segmentation scripts + auto layout#1201
Add AMP to ImageNet classification and segmentation scripts + auto layout#1201Kh4L wants to merge 2 commits intodmlc:masterfrom
Conversation
Kh4L
changed the title
|
Job PR-1201-3 is done. |
Kh4L
force-pushed
the
add_amp_classification
branch
from
21c9d60 to
52c5650
Compare
Kh4L
changed the title
Kh4L
force-pushed
the
add_amp_classification
branch
from
52c5650 to
a90357b
Compare
| assert not opt.auto_layout or opt.amp, "--auto-layout needs to be used with --amp" | ||
|
|
||
| if opt.amp: | ||
| amp.init(layout_optimization=opt.auto_layout) |
There was a problem hiding this comment.
Referring to definition of amp.init() here, seems there is no argument like layout_optimization ?
There was a problem hiding this comment.
It's an internal feature, it will be added soon
There was a problem hiding this comment.
Thanks for clarification.
There was a problem hiding this comment.
Just curiously, when setting both --amp and --dtype float16, what will be happening?
Kh4L
force-pushed
the
add_amp_classification
branch
2 times, most recently
from
5682d9a to
f51f405
Compare
|
Job PR-1201-9 is done. |
| help='whether to use group norm.') | ||
| parser.add_argument('--amp', action='store_true', | ||
| help='Use MXNet AMP for mixed precision training.') | ||
| parser.add_argument('--auto-layout', action='store_true', |
There was a problem hiding this comment.
Could you also add an option like --target-dtype since now we not only have float16 for amp, but bfloat16. Then, we can pass target-dtype to amp.init() to enable float16/bfloat16 training for GPU and CPU respectively. Thanks.
| if opt.resume_states is not '': | ||
| trainer.load_states(opt.resume_states) | ||
|
|
||
| if opt.amp: |
There was a problem hiding this comment.
Here may need change to if opt.amp and opt.target_dtype == 'float16':
| loss = [L(yhat, y.astype(opt.dtype, copy=False)) for yhat, y in zip(outputs, label)] | ||
| for l in loss: | ||
| l.backward() | ||
| if opt.amp: |
There was a problem hiding this comment.
Here may need change to if opt.amp and opt.target_dtype == 'float16':
|
|
||
| self.optimizer = gluon.Trainer(self.net.module.collect_params(), args.optimizer, | ||
| optimizer_params, kvstore=kv) | ||
| optimizer_params, update_on_kvstore=(False if args.amp else None)) |
There was a problem hiding this comment.
May I know why kvstore=kv is deleted? Could you add it back? Thanks.
| parser.add_argument('--syncbn', action='store_true', default=False, | ||
| help='using Synchronized Cross-GPU BatchNorm') | ||
| # performance related | ||
| parser.add_argument('--amp', action='store_true', |
There was a problem hiding this comment.
We usually add default=False for arguments. Could you add it? Thank you.
|
@Kh4L any update on this PR? |
|
@Kh4L Any update for this? BTW, do you have numbers for the improvement? |
Signed-off-by: Serge Panev <[email protected]>
Signed-off-by: Serge Panev <[email protected]>
Kh4L
force-pushed
the
add_amp_classification
branch
from
f51f405 to
f2e92a4
Compare
6 participants
Signed-off-by: Serge Panev [email protected]