NameError: name ‘custom_fwd’ is not defined

mixed-precision
1

Hi, I have a custom torch.autograd.Function function and try to use mix precision training but I got error below. I use pytorch 1.7

NameError: name ‘custom_fwd’ is not defined

Here is the example code

class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        pass
        return fwd_output
    @staticmethod
    @custom_bwd
    def backward(ctx, grad):
        pass
    
func=MyFloat32Func.apply
func(2)

And I got error

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-3-744176e249fa> in <module>
----> 1 class MyFloat32Func(torch.autograd.Function):
      2     @staticmethod
      3     @custom_fwd(cast_inputs=torch.float32)
      4     def forward(ctx, input):
      5         ctx.save_for_backward(input)

<ipython-input-3-744176e249fa> in MyFloat32Func()
      1 class MyFloat32Func(torch.autograd.Function):
      2     @staticmethod
----> 3     @custom_fwd(cast_inputs=torch.float32)
      4     def forward(ctx, input):
      5         ctx.save_for_backward(input)

NameError: name 'custom_fwd' is not defined

If I call the following print

print(torch.cuda.amp.custom_fwd)
print(torch.cuda.amp.custom_fwd(cast_inputs=torch.float32))

I get

<function custom_fwd at 0x7f9b0310cef0>
functools.partial(<function custom_fwd at 0x7f9b0310cef0>, cast_inputs=torch.float32)
2

You would have to either import custom_fwd from torch.cuda.amp or specify the full namespace as torch.cuda.amp.custom_fwd.

3

Thank you very much! @torch.cuda.amp.custom_fwd works!!

class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return torch.tensor(input,dtype=torch.float32)
    @staticmethod
    @torch.cuda.amp.custom_fwd
    def backward(ctx, grad):
        return 1
    
func=MyFloat32Func.apply
func(2)

Just wondering, why import custom_fwd from torch.cuda.amp does not work…

from torch.cuda.amp import custom_fwd


class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return torch.tensor(input,dtype=torch.float32)
    @staticmethod
    @custom_bwd
    def backward(ctx, grad):
        return 1
    
func=MyFloat32Func.apply
func(2)

Here is the error

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-12-d1461ca0ff6e> in <module>
      2 
      3 
----> 4 class MyFloat32Func(torch.autograd.Function):
      5     @staticmethod
      6     @custom_fwd(cast_inputs=torch.float32)

<ipython-input-12-d1461ca0ff6e> in MyFloat32Func()
      9         return torch.tensor(input,dtype=torch.float32)
     10     @staticmethod
---> 11     @custom_bwd
     12     def backward(ctx, grad):
     13         return 1

NameError: name 'custom_bwd' is not defined
4

You are only importing custom_fwd in the second approach, while you would need to import custom_fwd and custom_bwd.

Also, in your first example you are using custom_fwd twice, so you might want to change it as well.

5

Thank you very much!! It works fine now with custom_bwd above backward function!! :slight_smile: