实现简单的straight-through estimator(STE)（pytorch版）

 

import torch
input  = torch.randn(4,requires_grad = True)
output = torch.sign(input)
loss = output.mean()
loss.backward()
print(input)
print(input.grad)
 

tensor([-0.4608, 0.0240, 0.5585, -0.7694], requires_grad=True)
tensor([0., 0., 0., 0.])

import torch
from torch.autograd import Function
class LBSign(Function):
    @staticmethod
    def forward(self,input):
        return torch.sign(input)
    @staticmethod
    def backward(self,grad_output):
        return grad_output.clamp(-1,1)


sign = LBSign.apply
input = torch.randn(4,requires_grad = True)
output = sign(input)
loss = output.mean()
loss.backward()
print(input)
print(input.grad)


tensor([-1.0393, -0.5632,  0.5945, -0.5463], requires_grad=True)
tensor([0.2500, 0.2500, 0.2500, 0.2500])

参考链接：https://blog.csdn.net/ljs_a/article/details/81878466 

https://segmentfault.com/a/1190000020993594?utm_source=tag-newes