pytorch实现straightthroughestimator（STE）

现在深度学习中一般我们学习的参数都是连续的&＃xff0c;因为这样在反向传播的时候才可以对梯度进行更新。但是有的时候我们也会遇到参数是离>散的情况&＃xff0c;这样就没有办法进行反向传播了&＃xff0c;比如二值神经网络。本文中讲解了如何用pytorch对二值化的参数进行梯度更新的straight-through estimator算法。
Question&＃xff1a;
STE核心的思想就是我们的参数初始化的时候就是float这样的连续值&＃xff0c;当我们forward的时候就将原来的连续的参数映射到{-1, 1}带入到网络进行计算&＃xff0c;这样就可以计算网络的输出。然后backward的时候直接对原来float的参数进行更新&＃xff0c;而不是对二值化的参数更新。这样可以完成对整个网络的更新了。
首先我们对上面问题进行一下数学的讲解。

Example&＃xff1a;
首先我们验证一下使用torch.sign会是参数的梯度基本上都是0:

>>> input &＃61; torch.randn(4, requires_grad &＃61; True)
>>> output &＃61; torch.sign(input)
>>> loss &＃61; output.mean()
>>> loss.backward()
>>> input
tensor([-0.8673, -0.0299, -1.1434, -0.6172], requires_grad&＃61;True)
>>> input.grad
tensor([0., 0., 0., 0.])


 我们需要重写sign这个函数&＃xff0c;就好像写一个激活函数一样。

import torchclass LBSign(torch.autograd.Function):&＃64;staticmethoddef forward(ctx, input):return torch.sign(input)&＃64;staticmethoddef backward(ctx, grad_output):return grad_output.clamp_(-1, 1)


import torch
from LBSign import LBSignif __name__ &＃61;&＃61; &＃39;__main__&＃39;:sign &＃61; LBSign.applyparams &＃61; torch.randn(4, requires_grad &＃61; True) output &＃61; sign(params)loss &＃61; output.mean()loss.backward()


测试梯度&＃xff1a;

>>> params
tensor([-0.9143, 0.8993, -1.1235, -0.7928], requires_grad&＃61;True)
>>> params.grad
tensor([0.2500, 0.2500, 0.2500, 0.2500])