PyTorch自定义梯度

[3] 的一个简例。正向 sign：$$y=\text{sign}(x)=\{\begin{array}{llc}& -1& x<0\\ & 0& x=0\\ & 1& x>0\end{array}$$ 反向 Htanh：$$\frac{\partial y}{\partial x}=\{\begin{array}{llc}& 1& -ϵ\le x\le ϵ\\ & 0& else\end{array}$$

Code

import torch

class Htanh(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, epsilon=1):
        ctx.save_for_backward(x.data, torch.tensor(epsilon))
        return x.sign()

    @staticmethod
    def backward(ctx, dy):
        x, epsilon = ctx.saved_tensors
        dx = torch.where((x < - epsilon) | (x > epsilon), torch.zeros_like(dy), dy)
        return dx, None

• 用例

htanh = Htanh()
x = torch.tensor([-2.5, -1.6, 0, 3.6, 4.7])
y = htanh.apply(x)

References

1. PyTorch: Defining New autograd Functions
2. tensorflow自定义梯度
3. Binarized Neural Networks