Pytorch 1.1.0驾到！小升级大变动，易用性更强，支持自定义RNN

新智元

发布于 2019-05-13 20:04:24

1.1K0

发布于 2019-05-13 20:04:24

文章被收录于专栏：新智元新智元

【新智元导读】盼望已久，Pytorch终于更新了！Pytroch 1.1.0的发布除了修复了已有bug之外，最大的亮点就是可以更快、更好的支持自定义RNN，以及TensorBoard对可视化和模型调试提供了一流的本地支持。

Pytorch 1.1.0，来了！

可以说是一大波更新来袭了，话不多说上亮点：

TorchScript(Pytorch JIT)更快、更好的支持自定义RNN；
TensorBoard对可视化和模型调试提供了一流的本地支持；
可以在ScriptModule上通过使用torch.jit包装属性来分配属性；
TorchScript现在对列表和字典类型提供了鲁棒性的支持；
对于更复杂的有状态操作，TorchScript现在支持使用@torch.jit.script注释类；
nn.parallel.DistributedDataParallel：现在可以包装多GPU模块，它可以在一台服务器上实现模型并行和跨服务器的数据并行等用例。

注：不再支持CUDA 8.0。

此更新一出，在Reddit上也引发了一波热议，大部分网友们表示：

“赞！”、“好用！”、“爱了！”

用TorchScript优化CUDA递归神经网络

Pytorch添加的一个新特性是更好地支持带有TorchScript (PyTorch JIT)的快速自定义递归神经网络(fastrnns)。

RNN是一种流行的模型，在各种NLP任务上都表现出了良好的性能。PyTorch可以实现许多最流行的变体，例如Elman RNN、GRU和LSTM，以及多层和双向变体。

然而，许多用户希望实现他们自己的自定义RNN。将层规范化应用于LSTM就是这样一种用例。由于PyTorch CUDA LSTM实现使用融合内核，因此很难插入规范化甚至修改基本LSTM实现。

许多用户已经转向使用标准PyTorch运算符编写自定义实现，但是这样的代码遭受高开销：大多数PyTorch操作在GPU上启动至少一个内核，并且RNN由于其重复性质通常运行许多操作。但是可以应用TorchScript来融合操作并自动优化代码，在GPU上启动更少、更优化的内核。

此次更新的目标之一是让用户能够在TorchScript中编写快速，自定义的RNN，而无需编写专门的CUDA内核来实现类似的性能。接下来将提供如何使用TorchScript编写自己的快速RNN的教程。

编写自定义RNN

首先，可以使用下方链接中的文件作为模板来编写自己的自定义RNN。

https://github.com/pytorch/pytorch/blob/master/benchmarks/fastrnns/custom_lstms.py

如果想获得TorchScript当前提供的速度/优化（如运算符融合，批量矩阵乘法等），请遵循以下指南。

如果定制操作都是element-wise的，那就可以自动获得PyTorch JIT操作符fusion的优势！
如果有更复杂的操作（例如，reduce和element-wise的浑南和操作），请考虑分别对reduce操作和element-wise操作进行分组。
如果想知道自定义RNN中融合了什么，可以使用graph_for检查操作的优化图。以LSTMCell为例：

# get inputs and states for LSTMCell

 inputs = get_lstm_inputs()

 # instantiate a ScriptModule

 cell = LSTMCell(input_size, hidden_size)

 # print the optimized graph using graph_for

 out = cell(inputs)
 print(cell.graph_for(inputs))

这将提供的专用输入生成优化的TorchScript图形（a.k.a PyTorch JIT IR）：

graph(%x : Float(*, *),
         %hx : Float(*, *),
         %cx : Float(*, *),
         %w_ih : Float(*, *),
         %w_hh : Float(*, *),
         %b_ih : Float(*),
         %b_hh : Float(*)):
     %hy : Float(*, *), %cy : Float(*, *) = prim::DifferentiableGraph_0(%cx, %b_hh, %b_ih, %hx, %w_hh, %x, %w_ih)
     %30 : (Float(*, *), Float(*, *)) = prim::TupleConstruct(%hy, %cy)
     return (%30)
     with prim::DifferentiableGraph_0 = graph(%13 : Float(*, *),
         %29 : Float(*),
         %33 : Float(*),
         %40 : Float(*, *),
         %43 : Float(*, *),
         %45 : Float(*, *),
         %48 : Float(*, *)):
     %49 : Float(*, *) = aten::t(%48)
     %47 : Float(*, *) = aten::mm(%45, %49)
     %44 : Float(*, *) = aten::t(%43)
     %42 : Float(*, *) = aten::mm(%40, %44)
     ...some broadcast sizes operations...
     %hy : Float(*, *), %287 : Float(*, *), %cy : Float(*, *), %outgate.1 : Float(*, *), %cellgate.1 : Float(*, *), %forgetgate.1 : Float(*, *), %ingate.1 : Float(*, *) = prim::FusionGroup_0(%13, %346, %345, %344, %343)
     ...some broadcast sizes operations...
     return (%hy, %cy, %49, %44, %196, %199, %340, %192, %325, %185, %ingate.1, %forgetgate.1, %cellgate.1, %outgate.1, %395, %396, %287)
     with prim::FusionGroup_0 = graph(%13 : Float(*, *),
         %71 : Tensor,
         %76 : Tensor,
         %81 : Tensor,
         %86 : Tensor):
     ...some chunks, constants, and add operations...
     %ingate.1 : Float(*, *) = aten::sigmoid(%38)
     %forgetgate.1 : Float(*, *) = aten::sigmoid(%34)
     %cellgate.1 : Float(*, *) = aten::tanh(%30)
     %outgate.1 : Float(*, *) = aten::sigmoid(%26)
     %14 : Float(*, *) = aten::mul(%forgetgate.1, %13)
     %11 : Float(*, *) = aten::mul(%ingate.1, %cellgate.1)
     %cy : Float(*, *) = aten::add(%14, %11, %69)
     %4 : Float(*, *) = aten::tanh(%cy)
     %hy : Float(*, *) = aten::mul(%outgate.1, %4)
     return (%hy, %4, %cy, %outgate.1, %cellgate.1, %forgetgate.1, %ingate.1)

从上图中可以看到它有一个prim :: FusionGroup_0子图，它融合了LSTMCell中的所有element-wise操作（转置和矩阵乘法不是element-wise操作）。

可变长度序列最佳实践

TorchScript不支持PackedSequence。通常，当处理可变长度序列时，最好将它们填充到单个张量中并通过TorchScript LSTM发送该张量。例如：

sequences = [...] # List[Tensor], each Tensor is T' x C
padded = torch.utils.rnn.pad_sequence(sequences)
lengths = [seq.size(0) for seq in sequences]
padded  # T x N x C, where N is batch size and T is the max of all T'

model = LSTM(...)
output, hiddens = model(padded)
output  # T x N x C

当然，output可能在填充区域中有一些垃圾数据；使用lengths来跟踪你不需要的部分。

优化

现在将解释PyTorch JIT为加速自定义RNN所执行的优化。将在TorchScript中使用一个简单的自定义LSTM模型来说明优化，但其中许多是通用的并适用于其他RNN。

为了说明所做的优化以及如何从这些优化中获益，将运行一个用TorchScript编写的简单自定义LSTM模型（可以参考custom_lstm.py中的代码或下面的代码片段）并计算更改。

在配备2个Intel Xeon芯片和一个Nvidia P100的机器中设置环境，安装了cuDNN v7.3，CUDA 9.2。 LSTM模型的基本设置如下：

input_size = 512
hidden_size = 512
mini_batch = 64
numLayers = 1
seq_length = 100

PyTorch JIT最重要的是将python程序编译为PyTorch JIT IR，这是一个用于对程序图形结构进行建模的中间表示。然后，该IR可以从整个程序优化，硬件加速中受益，并且总体上具有提供大量计算增益的潜力。

接下来，将解释在如何提高训练或推理性能方面所做的主要优化，从LSTMCell和LSTMLayer开始，以及一些misc优化。

LSTM Cell(前向)

LSTM中的几乎所有计算都发生在LSTMCell中，因此重要的是看看它包含的计算以及如何提高它们的速度。下面是TorchScript中的LSTMCell实现示例：

class LSTMCell(jit.ScriptModule):
    def __init__(self, input_size, hidden_size):
        super(LSTMCell, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.weight_ih = Parameter(torch.randn(4 * hidden_size, input_size))
        self.weight_hh = Parameter(torch.randn(4 * hidden_size, hidden_size))
        self.bias_ih = Parameter(torch.randn(4 * hidden_size))
        self.bias_hh = Parameter(torch.randn(4 * hidden_size))

    @jit.script_method
    def forward(self, input, state):
        # type: (Tensor, Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tuple[Tensor, Tensor]]
        hx, cx = state
        gates = (torch.mm(input, self.weight_ih.t()) + self.bias_ih +
                 torch.mm(hx, self.weight_hh.t()) + self.bias_hh)
        ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)

        ingate = torch.sigmoid(ingate)
        forgetgate = torch.sigmoid(forgetgate)
        cellgate = torch.tanh(cellgate)
        outgate = torch.sigmoid(outgate)

        cy = (forgetgate * cx) + (ingate * cellgate)
        hy = outgate * torch.tanh(cy)

        return hy, (hy, cy)

TorchScript生成的此图形表示（IR）可实现多种优化和可伸缩计算。除了可以做的典型编译器优化（CSE，常量传播等）之外，还可以运行其他IR转换以使代码运行得更快。

LSTM层(前向)

class LSTMLayer(jit.ScriptModule):
    def __init__(self, cell, *cell_args):
        super(LSTMLayer, self).__init__()
        self.cell = cell(*cell_args)

    @jit.script_method
    def forward(self, input, state):
        # type: (Tensor, Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tuple[Tensor, Tensor]]
        inputs = input.unbind(0)
        outputs = torch.jit.annotate(List[Tensor], [])
        for i in range(len(inputs)):
            out, state = self.cell(inputs[i], state)
            outputs += [out]
        return torch.stack(outputs), state

在为TorchScript LSTM生成的IR上做了一些技巧来提高性能，团队做了一些示例优化：

循环展开(Loop Unrolling)：自动在代码中展开循环（对于大循环，展开它的一小部分），然后授权对for循环控制流进行进一步的优化。例如，fuser可以将循环体的迭代中的操作融合在一起，这导致对于诸如LSTM的控制流密集型模型的良好性能改进。
批量矩阵乘法：对于输入预乘的RNN（即模型具有大量相同LHS或RHS的矩阵乘法），可以将这些操作一起有效地批量处理为单个矩阵乘法，同时对输出进行分块以实现等效语义。

通过应用这些技术，将前向传播的时间减少了1.6ms，达到8.4ms(1.2倍加速)，后向传播的时间减少了7ms，达到20ms左右(1.35倍加速)。

LSTM层(后向)

“树结构”批处理矩阵Muplication：通常情况是在LSTM反向图中多次重复使用单个权重，形成一个树，其中叶子是矩阵乘法，节点是相加的。这些节点可以通过在不同维度上连接LHS和RHS来组合在一起，然后计算为单个矩阵乘法。等价公式可表示如下：

$L1 * R1 + L2 * R2 = torch.cat((L1, L2), dim=1) * torch.cat((R1, R2), dim=0)$

Autograd是使PyTorch成为如此优雅的ML框架的关键组件。因此，将其应用到PyTorch JIT，但是使用了一种新的自动微分(AD)机制，该机制在IR级别上工作。JIT自动微分将把正向图分割成符号可微分的子图，并为这些子图生成向后节点。以上面的IR为例，对于具有AD公式的操作，我们将图节点分组为一个prim :: DifferentiableGraph_0。对于没有添加到AD公式中的操作，我们将在执行期间返回到Autograd。
优化反向路径是困难的，隐式broadcasting语义使得自动微分的优化更加困难。 PyTorch可以方便地编写张量操作，而无需通过broadcasting张量来担心形状。对于性能而言，反向的痛点是需要对这种可broadcasting操作进行求和。这导致每个可broadcasting操作的导数后跟一个求和。由于目前无法融合减少操作，这会导致FusionGroups分成多个小组，从而导致性能下降。要解决这个问题，请参阅Thomas Viehmann撰写的文章：http://lernapparat.de/fast-lstm-pytorch/。

更多这方面的优化内容可参考Pytorch团队博客原文：

https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscript/

更多新功能

运算符