PyTorch自定义网络层
这篇博客关于如何通过自定义function和module来扩展 torch.nn和 torch.autograd,module对应于网络中的层。对于浅层网络,我们可以手写前向传播和后向传播过程。但是在深度学习中,网络结构很复杂,前向和后向也很复杂,手写就变得很困难。幸运的是,在PyTorch中存在自动微分的包,可以用来解决该问题。在使用自动微分包的时候,网络前向传播会定义一个计算图(computational graph),图中的节点是Tensor,两个节点之间的边对应了Tensor之间的变换关系函数。有了计算图的存在,Tensor的梯度计算也变得容易了些。例如,x是一个Tensor,其属性x.requires_grad = True,那么x.grad就是一个保存了这个Tensor x的梯度的一个标量值。
扩展 torch.autograd
最基础的自动求导操作在底层就是作用在两个张量上。前向传播函数是从输入张量到输出张量的计算过程;反向传播是输入输出张量的梯度(一些标量)并输出输入张量的梯度(一些标量)。在pytorch中我们可以很容易地定义自己的自动求导操作,通过继承torch.autograd.Function并定义forward和backward函数。
往autograd中添加操作需要给每一个操作实现一个新的Function子类。Function是autograd用来计算结果和梯度,保存操作记录的。每一个新的function都需要你去实现以下两个方法:
- forward():前向传播操作,它可以输入任意多的参数。任意的python对象都可以。在调用之前,记录了梯度的Tensor参数(如 requires_grad=True)会被转为不去记录梯度的参数。你既可以返回单一的Tensor,也可以返回Tensor的元组。
- backward():反向传播(梯度公式),输出的参数个数应与输入的参数个数一样多,每一个输出的参数代表着一个输入参数的梯度。如果你输入的参数不需要梯度(needs_input_grad是一个布尔值元组,表示是否需要梯度计算),或者输入不是Tensor对象,可以返回None.
下面是torch.nn里面的一个Linear function的代码:
# Inherit from Function
class LinearFunction(Function):
# Note that both forward and backward are @staticmethods
@staticmethod
# bias is an optional argument
def forward(ctx, input, weight, bias=None):
# ctx在这里类似self, ctx的属性可以在backward中调用
ctx.save_for_backward(input, weight, bias)
output = input.mm(weight.t())
if bias is not None:
output += bias.unsqueeze(0).expand_as(output)
return output
# This function has only a single output, so it gets only one gradient
@staticmethod
def backward(ctx, grad_output):
# This is a pattern that is very convenient - at the top of backward
# unpack saved_tensors and initialize all gradients w.r.t. inputs to
# None. Thanks to the fact that additional trailing Nones are
# ignored, the return statement is simple even when the function has
# optional inputs.
# 在这里,获取ctx中保存的参数
input, weight, bias = ctx.saved_tensors
grad_input = grad_weight = grad_bias = None
# These needs_input_grad checks are optional and there only to
# improve efficiency. If you want to make your code simpler, you can
# skip them. Returning gradients for inputs that don't require it is
# not an error.
if ctx.needs_input_grad[0]:
grad_input = grad_output.mm(weight)
if ctx.needs_input_grad[1]:
grad_weight = grad_output.t().mm(input)
if bias is not None and ctx.needs_input_grad[2]:
grad_bias = grad_output.sum(0).squeeze(0)
return grad_input, grad_weight, grad_bias
为了使用简单些,我们建议给apply 方法换个名字:
linear = LinearFunction.apply
这里,给了另一个例子,参数不是Tensor,而是常量,不需要更新:
class MulConstant(Function):
@staticmethod
def forward(ctx, tensor, constant):
# ctx is a context object that can be used to stash information
# for backward computation
ctx.constant = constant
return tensor * constant
@staticmethod
def backward(ctx, grad_output):
# We return as many input gradients as there were arguments.
# Gradients of non-Tensor arguments to forward must be None.
return grad_output * ctx.constant, None
你可能想检查下backward方法计算function的梯度是否准确,你可以和数值近似方法得到的结果来比较:
from torch.autograd import gradcheck
# gradcheck takes a tuple of tensors as input, check if your gradient
# evaluated with these tensors are close enough to numerical
# approximations and returns True if they all verify this condition.
input = (torch.randn(20,20,dtype=torch.double,requires_grad=True), torch.randn(30,20,dtype=torch.double,requires_grad=True))
test = gradcheck(linear, input, eps=1e-6, atol=1e-4)
print(test)
扩展torch.nn
计算图和自动微分在定义复杂网络和求梯度的时候非常好用,但对于大型网络,这个有点太底层。在我们构建网络的时候,经常希望将计算限制在每个层里面(参数分层更新)。在PyTorch中提供了nn包,定义了一组等价于layer的模块module。一个module接受输入tensor并得到输出tensor,同时也会包含可学习的参数。
有时候我们需要运用一些新的且nn中没有的module,此时就需要自定义自己的module了,自定义的module需要继承nn.Module且自定义forward函数。其中forward函数可以接受输入tensor并利用其他模型或者其他自动给求导操作来产生输出tensor. 但并不需要重写backward函数,因此nn使用了autograd. 这也就意味着,需要自定义module, 都必须要有对应的autograd函数调用其中的backward。
增加一个Module
因为nn经常使用autograd,增加一个新的Module需要实现一个计算梯度的function。例如,我们想实现一个Linear module,而且我们已经在上面实现了它的function。我们要实现两个function:
- init(optional):需要如kernel_size, 特征个数等参数,而且初始化参数和buffer。
- forward():实例化一个function并使用它来执行操作。这跟上面的function wrapper很类似。
下面就是Linear module如何实现的:
class Linear(nn.Module):
def __init__(self, input_features, output_features, bias=True):
super(Linear, self).__init__()
self.input_features = input_features
self.output_features = output_features
# nn.Parameter is a special kind of Tensor, that will get
# automatically registered as Module's parameter once it's assigned
# as an attribute. Parameters and buffers need to be registered, or
# they won't appear in .parameters() (doesn't apply to buffers), and
# won't be converted when e.g. .cuda() is called. You can use
# .register_buffer() to register buffers.
# nn.Parameters require gradients by default.
self.weight = nn.Parameter(torch.Tensor(output_features, input_features))
if bias:
self.bias = nn.Parameter(torch.Tensor(output_features))
else:
# You should always register all possible parameters, but the
# optional ones can be None if you want.
self.register_parameter('bias', None)
# Not a very smart way to initialize weights
self.weight.data.uniform_(-0.1, 0.1)
if bias is not None:
self.bias.data.uniform_(-0.1, 0.1)
def forward(self, input):
# See the autograd section for explanation of what happens here.
return LinearFunction.apply(input, self.weight, self.bias)
def extra_repr(self):
# (Optional)Set the extra information about this module. You can test
# it by printing an object of this class.
return 'in_features={}, out_features={}, bias={}'.format(
self.in_features, self.out_features, self.bias is not None
)
Function与Module异同:
Function与Module都可以对pytorch进行自定义拓展,使其满足网络的需求,但这两者还是有十分重要的不同:
- Function一般只定义一个操作,因为其无法保存参数,因此适用于激活函数、pooling等操作;Module是保存了参数,因此适合于定义一层,如线性层,卷积层,也适用于定义一个网络;
- Function需要定义三个方法:init, forward, backward(需要自己写求导公式);Module:只需定义init和forward,而backward的计算由自动求导机制构成;
- 可以不严谨的认为,Module是由一系列Function组成,因此其在forward的过程中,Function和Variable组成了计算图,在backward时,只需调用Function的backward就得到结果,因此Module不需要再定义backward;
- Module不仅包括了Function,还包括了对应的参数,以及其他函数与变量,这是Function所不具备的;
- Module 是PyTorch组织神经网络的基本方式。Module 包含了模型的参数以及计算逻辑。Function承载了实际的功能,定义了前向和后向的计算逻辑;
- Module 是任何神经网络的基类,PyTorch中所有模型都必需是 Module的子类。Module 可以套嵌,构成树状结构。一个 Module 可以通过将其他 Module 做为属性的方式,完成套嵌。
- Function 是 PyTorch自动求导机制的核心类。Function是无参数或者说无状态的,它只负责接收输入,返回相应的输出;对于反向,它接收输出相应的梯度,返回输入相应的梯度。
- 在调用loss.backward()时,使用的是Function子类中定义的backward()函数。