[create_supervised_trainer] add automatic mixed precision

[ydcjeff]

Fixes #1235

Description: Add automatic mixed precision using torch.cuda.amp and apex.

Usage:

# using autocast only
trainer = create_supervised_trainer(amp_mode='amp')

# using autocast and default scaler
trainer = create_supervised_trainer(amp_mode='amp', scaler=True)
# trainer state will have attribute scaler
print(trainer.state.scaler)
<torch.cuda.amp.grad_scaler.GradScaler object at 0x7f8e0dac7b80>

# using autocast and custom scaler
# trainer state will not have attribute scaler if scaler instance is passed
scaler = GradScaler(2**10)
trainer = create_supervised_trainer(amp_mode='amp', scaler=scaler)

# using apex
trainer = create_supervised_trainer(amp_mode='apex')

# scaler will be ignored, warning will show up
trainer = create_supervised_trainer(amp_mode='apex', scaler=True)

Check list:

• New tests are added (if a new feature is added)
• New doc strings: description and/or example code are in RST format
• Documentation is updated (if required)

https://deploy-preview-1589--pytorch-ignite-preview.netlify.app/engine.html#

[vfdev-5]

Thanks a lot for the PR @ydcjeff !

[vfdev-5]

Currently, launched it manually : https://app.circleci.com/pipelines/github/pytorch/ignite/1195/workflows/27d5b840-72bb-41e1-8d1c-84640f1f623c but I think either next commit or new PR will run auto on GPUs

[ydcjeff]

Currently, launched it manually : https://app.circleci.com/pipelines/github/pytorch/ignite/1195/workflows/27d5b840-72bb-41e1-8d1c-84640f1f623c but I think either next commit or new PR will run auto on GPUs

Thank you!

[ydcjeff]

I need help with the tests specifically with apex and GradScaler.

[vfdev-5]

I need help with the tests specifically with apex and GradScaler.

I'll try to implement something from my side and we'll see.

[vfdev-5]

We discussed this PR and related issue with the team and we think that we should explore a bit different approach. Helper method create_supervised_trainer is roughly made of 2 things : update function definition and Engine setup.

Probably, it would more helpful to provide public methods like:

• supervised_training_step
• supervised_training_step_tpu
• supervised_training_step_apex
• supervised_training_step_amp

and inside create_supervised_trainer we could setup the trainer according to provided options without lots of if/else. Maybe, we can skip for instance grad norm.

Basically, the idea is something like that :

def get_training_step_1(a):
    def training_step(e, b):
        print(a, e, b)
    return training_step
    
def get_training_step_2(a):
    def training_step(e, b):
        print(a, e, b, "with amp")
    return training_step

def create_supervised_trainer(a, opt):
    training_step = None
    if opt == 1:
        training_step = get_training_step_1(a)
    elif opt == 2:
        training_step = get_training_step_2(a)
        
    e = Engine(training_step)
    return e  

cc @sdesrozis any other ideas or thoughts ?

[sdesrozis]

That would be great for users to have these functions, helpful to check under the hood.

My thoughts on this topic is about the update function. The dream would be to pass a generic function and have an automatic (or close) tools to adapt to features like amp, tpu, etc.

[ydcjeff]

Shall we also accept scaler argument or create only internally?

[vfdev-5]

Thanks for the update!

[vfdev-5]

@sdesrozis can you review the PR please

[ydcjeff]

Thank you @sdesrozis for your review.
if gpu tests pass, we are ready to merge.

[vfdev-5]

Looks good to me as well ! Thanks a lot @ydcjeff !
I left few nit comments about removing TPU mentions where it is inappropriate.
The comment about the warning and adding usage examples of these features could be done in a follow-up PR...

[sdesrozis]

Could we add tests to decrease codecov warnings ?

[vfdev-5]

Could we add tests to decrease codecov warnings ?

Let's do that all in a follow-up PR :)

[ydcjeff]

Could we add tests to decrease codecov warnings ?

I think those warnings are from one_gpu_tests failing to upload coverage.

The comment about the warning and adding usage examples of these features could be done in a follow-up PR...

will do that.

[ydcjeff]

Found out that amp module is available in torch 1.5 which doesn't have autocast yet.
Changed to ImportError to handle all torch version.

[sdesrozis]

I think those warnings are from one_gpu_tests failing to upload coverage.

Do we know why it does not work ?

[ydcjeff]

I think those warnings are from one_gpu_tests failing to upload coverage.

Do we know why it does not work ?

I don't know very well, but it can be codecov failed to upload to its server.

[vfdev-5]

I think those warnings are from one_gpu_tests failing to upload coverage.

Do we know why it does not work ?

Asked here: codecov/codecov-bash#411

Anyway, if there is no way to fix the uploading we can remove - Z option and silently ignore uploading issue.

[vfdev-5]

@sdesrozis can you please merge this PR once the ci is OK.

[ydcjeff]

Thank you for your help and reviews.