Pytorch 学习笔记之自定义 Module

pytorch 是一个基于 python 的深度学习库。pytorch 源码库的抽象层次少,结构清晰,代码量适中。相比于非常工程化的 tensorflow,pytorch 是一个更易入手的,非常棒的深度学习框架。

对于系统学习 pytorch,官方提供了非常好的入门教程 ,同时还提供了面向深度学习的示例,同时热心网友分享了更简洁的示例

1. overview

不同于 theano,tensorflow 等低层程序库,或者 keras、sonnet 等高层 wrapper,pytorch 是一种自成体系的深度学习库(图1)。

图1. 几种深度学习程序库对比

如图2所示,pytorch 由低层到上层主要有三大块功能模块。

图2. pytorch 主要功能模块

1.1 张量计算引擎(tensor computation)

Tensor 计算引擎,类似 numpy 和 matlab,基本对象是tensor(类比 numpy 中的 ndarray 或 matlab 中的 array)。除提供基于 CPU 的常用操作的实现外,pytorch 还提供了高效的 GPU 实现,这对于深度学习至关重要。

1.2 自动求导机制(autograd)

由于深度学习模型日趋复杂,因此,对自动求导的支持对于学习框架变得必不可少。pytorch 采用了动态求导机制,使用类似方法的框架包括: chainer,dynet。作为对比,theano,tensorflow 采用静态自动求导机制。

1.3 神经网络的高层库(NN)

pytorch 还提供了高层的神经网络模块。对于常用的网络结构,如全连接、卷积、RNN 等。同时,pytorch 还提供了常用的目标函数optimizer 及参数初始化方法

这里,我们重点关注如何自定义神经网络结构。

2. 自定义 Module

图3. pytorch Module

module 是 pytorch 组织神经网络的基本方式。Module 包含了模型的参数以及计算逻辑。Function 承载了实际的功能,定义了前向和后向的计算逻辑。

下面以最简单的 MLP 网络结构为例,介绍下如何实现自定义网络结构。完整代码可以参见repo

2.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 中反复调用。

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 是自动求导机制的核心类,此处暂不介绍,参见教程

3. 自定义循环神经网络(RNN)

我们尝试自己定义一个更复杂的 Module ——RNN。这里,我们只定义最基础的 vanilla RNN(图4),基本的计算公式如下:

ht=relu(W⋅x+U⋅ht−1)

图4. RNN【来源】

更复杂的 LSTM、GRU 或者其他变种的实现也非常类似。

3.1 定义 Cell

import torch
from torch.nn import Module, Parameter

class RNNCell(Module):
    def __init__(self, input_size, hidden_size):
        super(RNNCell, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size

        self.weight_ih = Parameter(torch.Tensor(hidden_size, input_size))
        self.weight_hh = Parameter(torch.Tensor(hidden_size, hidden_size))
        self.bias_ih = Parameter(torch.Tensor(hidden_size))
        self.bias_hh = Parameter(torch.Tensor(hidden_size))

        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1.0 / math.sqrt(self.hidden_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, stdv)

    def forward(self, input, h):
        output = LinearF()(input, self.weight_ih, self.bias_ih) + LinearF()(h, self.weight_hh, self.bias_hh)
        output = ReLUF()(output)

        return output

3.2 定义完整的 RNN

import torch
from torch.nn import Module

class RNN(Moudule):
    def __init__(self, input_size, hidden_size):
        super(RNN, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size

        sef.cell = RNNCell(input_size, hidden_size)

    def forward(self, inputs, initial_state):
        time_steps = inputs.size(1)

        state = initial_state
        outputs = []
        for t in range(time_steps):
            state = self.cell(inputs[:, t, :], state)
            outputs.append(state)

        return outputs

可运行的完整代码见repo

讨论

pytorch 的 Module 结构是传承自 torch,这一点也同样被 keras (functional API)所借鉴。 在 caffe 等一些[早期的]深度学习框架中,network 是由于若干 layer ,经由不同的拓扑结构组成的。而在 (pyt)torch 中没有 layer 和 network 是区分,一切都是 callable 的 Module。Module 的调用的输入和输出都是 tensor (由 Variable 封装),用户可以非常自然的构造任意有向无环的网络结构(DAG)。

同时, pytorch 的 autograd 机制封装的比较浅,可以比较容易的定制反传或修改梯度。这对有些算法是非常重要。

总之,仅就自定义算法而言,pytorch 是一个非常优雅的深度学习框架。

原创声明,本文系作者授权云+社区-专栏发表,未经许可,不得转载。

如有侵权,请联系 yunjia_community@tencent.com 删除。

编辑于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏AI科技大本营的专栏

入门 | 海量数据处理算法总结【超详解】

作者 | Angel_Kitty ➤1. Bloom Filter 【Bloom Filter】 Bloom Filter(BF)是一种空间效率很高的随机数据...

3679
来自专栏灯塔大数据

每周学点大数据 | No.23 外排序(二)

No.23期 外排序(二) Mr. 王:接下来我们用一个例子对磁盘归并排序进行说明。先来约定讨论的参数:N=24,M=8,B=2。 小可:嗯,一共有2...

3226
来自专栏技术总结

算法(3)

1887
来自专栏杂七杂八

numpy科学计算包的使用2

利用数组进行数据处理 NumPy数组使你可以将许多种数据处理任务表述为简洁的数组表达式(否则需要编写循环)。用数组表达式代替循环的做法,通常被称为矢量化。 矢...

34712
来自专栏竹清助手

【机器学习】 搭建模型第一步:你需要预习的NumPy基础都在这了

NumPy 主要的运算对象为同质的多维数组,即由同一类型元素(一般是数字)组成的表格,且所有元素通过正整数元组进行索引。在 NumPy 中,维度 (dimens...

1024
来自专栏xcywt

《大话数据结构》一些基础知识

第一章 数据结构绪论 1.4 基本概念和术语 1.4.1 数据 数据:描述客观事物的符号,是计算机中可以操作的对象,是能被极端及识别,并输入给计算机处理的符号集...

3349
来自专栏java技术学习之道

海量数据处理 - 找出最大的n个数(top K问题)

1384

与机器学习算法有关的数据结构

可能你对经常使用的统计分类包中的功能不满足你的需求而感到不爽,或者你已经有了一个新的数据处理方法。所以,你决定改动现有封装好的算法,开始编写你自己的机器学习方法...

2347
来自专栏Petrichor的专栏

numpy: 常用api速查

1077
来自专栏简书专栏

基于Numpy的线性代数运算

numpy.matrix方法的参数可以为ndarray对象 numpy.matrix方法的参数也可以为字符串str,示例如下:

1003

扫码关注云+社区