How to compute a block of code using full precision when training with AMP?

[GuillaumeTong]

I have encountered a few cases so far where some losses (such as the ms-ssim implementation from VainF on GitHub ) that used to work under full precision training will no longer work when using AMP.
I would like to know if there is any way to locally disable AMP for just a few steps of computation, either through Lightning or through native PyTorch. (It seems that torch.cuda.amp.autocast(enabled=False) might be able to do the job, but I am worried about it messing with the Lightning internals)
Since I calculate a variety of losses before combining them, all in the same step, I am thinking that doing so will still allow me to leverage some of the advantages from AMP

[stale[bot]]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[GuillaumeTong]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

Nobody even replied...

[carmocca]

Hi @GuillaumeTong!

As you mentioned, you should be able to disable AMP for a region with PyTorch's autocast context manager.

Sorry I don't have a clear answer for you, but trying it out will be your best bet to know if it will work, as it also depends on your particular use case.

If you find any bugs, feel free to open an issue and we will try to help you asap 😄

[tchaton]

Hey @GuillaumeTong,

Context manager can be nested.

Can you try the following:

Trainer(amp_backend="native", precision=16)

In your code.

@autocast(False)
def compute_loss

Best,
T.C