pytorch自定义Module,并通过其理解DoReFaNet网络定义方法。
参考中文官方,详情参考：PyTorch 如何自定义 Module
1.自定义Module
Module 是 pytorch 组织神经网络的基本方式。Module 包含了模型的参数以及计算逻辑。Function 承载了实际的功能，定义了前向和后向的计算逻辑。
下面以最简单的 MLP 网络结构为例，介绍下如何实现自定义网络结构。完整代码可以参见repo。
1.1 Function
Function 是 pytorch 自动求导机制的核心类。Function 是无参数或者说无状态的，它只负责接收输入，返回相应的输出；对于反向，它接收输出相应的梯度，返回输入相应的梯度。
这里我们只关注如何自定义 Function。Function 的定义见源码。下面是简化的代码段：
class Function(object):
def forward(self, *input):
raise NotImplementedError
def backward(self, *grad_output):
raise NotImplementedError
forward 和 backward 的输入和输出都是 Tensor 对象。
Function 对象是 callable 的，即可以通过()的方式进行调用。其中调用的输入和输出都为 Variable 对象。下面的代码示例了如何实现一个 ReLU **函数并进行调用：
import torch
from torch.autograd import Function
class ReLUF(Function)：
def forward(self, input):
self.save_for_backward(input)
output = input.clamp(min=0)
return output
def backward(self, output_grad):
input = self.to_save[0]
input_grad = output_grad.clone()
input_grad[input < 0] = 0
return input_grad
## Test
if __name__ == "__main__":
from torch.autograd import Variable
torch.manual_seed(1111)
a = torch.randn(2, 3)
va = Variable(a, requires_grad=True)
vb = ReLUF()(va)
print va.data, vb.data
vb.backward(torch.ones(va.size()))
print vb.grad.data, va.grad.data
如果 backward 中需要用到 forward 的输入，需要在 forward 中显式的保存需要的输入。在上面的代码中，forward 利用self.save_for_backward函数，将输入暂时保存，并在 backward 中利用saved_tensors (python tuple 对象) 取出。
显然，forward 的输入应该和 backward 的输入相对应；同时，forward 的输出应该和 backward 的输入相匹配。
注意：
由于 Function 可能需要暂存 input tensor，因此，建议不复用 Function 对象，以避免遇到内存提前释放的问题。
如示例代码所示，forward的每次调用都重新生成一个 ReLUF 对象，而不能在初始化时生成在 forward 中反复调用。（意思示例代码是对的，其是每个fc都生成一个新对象。因为每个对象都自己在管理自己的的权重跟输入）
2.2 Module
类似于 Function，Module 对象也是 callable 是，输入和输出也是 Variable。不同的是，Module 是[可以]有参数的。Module 包含两个主要部分：参数及计算逻辑（Function 调用）。由于ReLU**函数没有参数，这里我们以最基本的全连接层为例来说明如何自定义Module。
全连接层的运算逻辑定义如下 Function:
import torch
from torch.autograd import Function
class LinearF(Function):
def forward(self, input, weight, bias=None):
self.save_for_backward(input, weight, bias)
output = torch.mm(input, weight.t())
if bias is not None:
output += bias.unsqueeze(0).expand_as(output)
return output
def backward(self, grad_output):
input, weight, bias = self.saved_tensors
grad_input = grad_weight = grad_bias = None
if self.needs_input_grad[0]:
grad_input = torch.mm(grad_output, weight)
if self.needs_input_grad[1]:
grad_weight = torch.mm(grad_output.t(), input)
if bias is not None and self.needs_input_grad[2]:
grad_bias = grad_output.sum(0).squeeze(0)
if bias is not None:
return grad_input, grad_weight, grad_bias
else:
return grad_input, grad_weight
needs_input_grad 为一个元素为 bool 型的 tuple，长度与 forward 的参数数量相同，用来标识各个输入是否输入计算梯度；对于无需梯度的输入，可以减少不必要的计算。
Function（此处为 LinearF） 定义了基本的计算逻辑，Module 只需要在初始化时为参数分配内存空间，并在计算时，将参数传递给相应的 Function 对象。代码如下：
import torch
import torch.nn as nn
class Linear(nn.Module):
def __init__(self, in_features, out_features, bias=True):
super(Linear, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
if bias:
self.bias = nn.Parameter(torch.Tensor(out_features))
else:
self.register_parameter('bias', None)
def forward(self, input):
return LinearF()(input, self.weight, self.bias)
需要注意的是，参数是内存空间由 tensor 对象维护，但 tensor 需要包装为一个Parameter 对象。Parameter 是 Variable 的特殊子类，仅有是不同是 Parameter 默认requires_grad为 True。Varaible 是自动求导机制的核心类，此处暂不介绍，参见教程。
2.3自定义循环神经网络（RNN）
可运行代码参考：
RNN
其中Parameters是Variable的一个子类，而且其是自动求导机制，所以我们在定义代码网络的时候基本不需要从在backwards，只要定义好forward就可以了
2.3 定义DoReFaNet
由于其压缩的是权重，所以我们只需要定义一个卷积类，**函数类，然后去获取他里面的权重，进行量化，定义网络如下：
class AlexNet_Q(nn.Module):
def __init__(self, wbit, abit, num_classes=1000):
super(AlexNet_Q, self).__init__()
Conv2d = conv2d_Q_fn(w_bit=wbit)
Linear = linear_Q_fn(w_bit=wbit)
self.features = nn.Sequential(
nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=2),
nn.BatchNorm2d(96),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
Conv2d(96, 256, kernel_size=5, padding=2),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
activation_quantize_fn(a_bit=abit),
nn.MaxPool2d(kernel_size=3, stride=2),
Conv2d(256, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
activation_quantize_fn(a_bit=abit),
Conv2d(384, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
activation_quantize_fn(a_bit=abit),
Conv2d(384, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
activation_quantize_fn(a_bit=abit),
nn.MaxPool2d(kernel_size=3, stride=2),
)
self.classifier = nn.Sequential(
Linear(256 * 6 * 6, 4096),
nn.ReLU(inplace=True),
activation_quantize_fn(a_bit=abit),
Linear(4096, 4096),
nn.ReLU(inplace=True),
activation_quantize_fn(a_bit=abit),
nn.Linear(4096, num_classes),
)
for m in self.modules():
if isinstance(m, Conv2d) or isinstance(m, Linear):
init.xavier_normal_(m.weight.data)
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), 256 * 6 * 6)
x = self.classifier(x)
return
其中自定义权重量化如下：
class weight_quantize_fn(nn.Module):
def __init__(self, w_bit):
super(weight_quantize_fn, self).__init__()
assert w_bit <= 8 or w_bit == 32
self.w_bit = w_bit
self.uniform_q = uniform_quantize(k=w_bit)
def forward(self, x):
if self.w_bit == 32:
weight_q = x
elif self.w_bit == 1:
E = torch.mean(torch.abs(x)).detach()
weight_q = self.uniform_q(x / E) * E
else:
weight = torch.tanh(x)
weight = weight / 2 / torch.max(torch.abs(weight)) + 0.5
weight_q = 2 * self.uniform_q(weight) - 1
return weight_q
def conv2d_Q_fn(w_bit):
class Conv2d_Q(nn.Conv2d):
def __init__(self, in_channels, out_channels, kernel_size, stride=1,
padding=0, dilation=1, groups=1, bias=True):
super(Conv2d_Q, self).__init__(in_channels, out_channels, kernel_size, stride,
padding, dilation, groups, bias)
self.w_bit = w_bit
self.quantize_fn = weight_quantize_fn(w_bit=w_bit)
def forward(self, input, order=None):
weight_q = self.quantize_fn(self.weight)
# print(np.unique(weight_q.detach().numpy()))
return F.conv2d(input, weight_q, self.bias, self.stride,
self.padding, self.dilation, self.groups)
return Conv2d_Q