pytorch实现直通估计器（STE）

原理部分参考链接：
https://segmentfault.com/a/1190000020993594?utm_source=tag-newest

在此就不对原理做过多阐述了，原文中写的非常好，简而言之就是反向传播过程中应用Htanh函数的梯度结果替换torch.sign的恒为零的梯度。

我对链接作者代码的反向传播部分存有疑惑，.clamp_(-1, 1)函数只是将后一层（即距离输出层较近的那一层）的梯度做了截断，但Htanh函数的梯度和输入有关，但与后一层的梯度无关，也就是说这并不是计算Htanh的梯度，所以对后一层的梯度做截断我认为并不正确。

个人认为正确的做法应该是当输入的绝对值小于1时，直接传递后一层的梯度；当输入的绝对值大于1时，使其梯度归零。
以下为pytorch代码：

#STE
import torch

class SignSTE(torch.autograd.Function):
    @staticmethod                         # 静态方法，不需要实例化
    def forward(ctx, input):
        ctx.save_for_backward(input)      # 保存输入以备在backward中使用
        return torch.sign(input)
    
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors        # 获取保存的输入张量
        grad_Htanh = grad_output.clone()  # 复制一份grad_output，以便我们可以修改它
        grad_Htanh[input.abs() > 1] = 0   # 当输入绝对值大于1时，将梯度置为0
        return grad_Htanh

参考文章：
[1]: https://segmentfault.com/a/1190000020993594?utm_source=tag-newest
[2]: https://blog.csdn.net/m0_37400316/article/details/105996240