文章目錄
  1. 1. Question
  2. 2. Example
    1. 2.1. torch.sign
    2. 2.2. demo
    3. 2.3. explain
  3. 3. reference

现在深度学习中一般我们学习的参数都是连续的,因为这样在反向传播的时候才可以对梯度进行更新。但是有的时候我们也会遇到参数是离散的情况,这样就没有办法进行反向传播了,比如二值神经网络。本文中讲解了如何用pytorch对二值化的参数进行梯度更新。

Question

提示: 保证公式能正常显示请使用google chrome
STE核心的思想就是我们的参数初始化的时候就是float这样的连续值,当我们forward的时候就将原来的连续的参数映射到{-1,, 1}带入到网络进行计算,这样就可以计算网络的输出。然后backward的时候直接对原来float的参数进行更新,而不是对二值化的参数更新。这样可以完成对整个网络的更新了。
首先我们对上面问题进行一下数学的讲解。

  • 我们希望参数的范围是$r \in \mathbb{R}$
  • 我们可以得到二值化的参数 $q = Sign(r)$, $Sign$函数可以参考torch.sign函数, 可以理解为取符号函数
  • backward的过程中对$q$求梯度可得 $\frac{\partial loss}{\partial q}$
  • 对于$\frac{\partial q}{\partial r} = 0$, 所以可以得出 $\frac{\partial loss}{\partial r} = 0$, 这样的话我们就无法完成对参数的更新,因为每次lossr梯度都是0
  • 所以backward的过程我们需要修改$\frac{\partial q}{\partial r}$这部分才可以使梯度继续更新下去,所以对$\frac{\partial loss}{\partial r}$进行如下修改: $\frac{\partial q}{\partial r} = \frac{\partial loss}{\partial q} * 1_{|r| \leq 1}$, 其中
    $1_{|r| \leq 1}$ 可以看作$Htanh(x) = Clip(x, -1, 1) = max(-1, min(1, x))$对$x$的求导过程, 也就是是说:
    $$\frac{\partial loss}{\partial r} = \frac{\partial loss}{\partial q} \frac{\partial Htanh}{\partial r}$$

Example

torch.sign

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

1
2
3
4
5
6
7
8
>>> input = torch.randn(4, requires_grad = True)
>>> output = torch.sign(input)
>>> loss = output.mean()
>>> loss.backward()
>>> input
tensor([-0.8673, -0.0299, -1.1434, -0.6172], requires_grad=True)
>>> input.grad
tensor([0., 0., 0., 0.])

demo

我们需要重写sign这个函数,就好像写一个激活函数一样。先看一下代码, github源码:
LBSign.py

1
2
3
4
5
6
7
8
9
10
11
import torch
class LBSign(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
return torch.sign(input)
@staticmethod
def backward(ctx, grad_output):
return grad_output.clamp_(-1, 1)

接下来我们做一下测试
main.py

1
2
3
4
5
6
7
8
9
10
import torch
from LBSign import LBSign
if __name__ == '__main__':
sign = LBSign.apply
params = torch.randn(4, requires_grad = True)
output = sign(params)
loss = output.mean()
loss.backward()

然后我们发现有梯度了

1
2
3
4
>>> params
tensor([-0.9143, 0.8993, -1.1235, -0.7928], requires_grad=True)
>>> params.grad
tensor([0.2500, 0.2500, 0.2500, 0.2500])

explain

接下来我们对代码就行一下解释pytorch文档链接:

  • forward中的参数ctx是保存的上下文信息,input是输入
  • backward中的参数ctx是保存的上下文信息,grad_output可以理解成 $\frac{\partial loss}{\partial q}$这一步的梯度信息,我们需要做的就是让
    $$gradoutput * \frac{\partial Htanh}{\partial r}$$ 而不是让pytorch继续默认的 $$gradoutput * \frac{\partial q}{\partial r}$$
    但是我们可以从上面的公式可以看出函数$Htanh$对$x$求导是1, 当$x \in [-1, 1]$,所以程序就可以化简成保留原来的梯度就行了,然后裁剪到其他范围的。

reference

torch.autograd.Function

Binarized Neural Networks: Training Deep Neural Networks with Weights and Activations Constrained to +1 or -1

二值网络,围绕STE的那些事儿

Custom binarization layer with straight through estimator gives error

定义torch.autograd.Function的子类,自己定义某些操作,且定义反向求导函数

文章目錄
  1. 1. Question
  2. 2. Example
    1. 2.1. torch.sign
    2. 2.2. demo
    3. 2.3. explain
  3. 3. reference