Disabling mixed precsion in my own layers

Hi, after reading the docs about mixed precsion, amp_example
I’m still confused with several problems.

Let’s say if I have two networks, one is the standard resnet50 and another is a sparse conv layer.
input images are first passed through resnet50 and then sparse convs.

If I only want to use half for resnet and keep float32 for the sparse conv layer (so I don’t have to modify the code)
I only need to warp the model within the autocast function and disable it before the sparse conv layers?
like,

with autocast():
      out = resnet50(x)
      with autocast(enabled=False):
             out = sparseconv(out.float())

right?

And from my knowledge, gradients are scaled during mixed precision,
If I have to write my own backward function for sparse conv layers (warped in autocast(disabled)),
Do I still need to consider the scale?

Your autocast example looks correct.

Yes, if you are using a GradScaler object.

No, gradient scaling and auto casting are working together but are independent from each other.
You don’t need to explicitly implement gradient scaling for the custom layer.

1 Like