开发 | 从原理到实战 英伟达教你用PyTorch搭建RNN（下）

代码实操

在开始创建神经网络之前，我需要设置一个 data loader。对于深度学习而言，在数据样例的 batch 上运行模型十分常见，这能通过并行计算加速训练，并在每一步有更加平滑的梯度。现在我们就开始，下文会解释上篇描述的如何对 stack-manipulation 进行 batch。 PyTorch text library 内置的系统，能把相近长度的样例组合起来自动生成 batch，以下 Python 代码便向该系统加载了一些数据。运行这些代码之后，, train_iter、dev_iter、test_iter 中的迭代器，会在 SNLI 训练、验证、测试阶段在 batch 上循环。

from torchtext import data, datasets
TEXT = datasets.snli.ParsedTextField(lower=True)
TRANSITIONS = datasets.snli.ShiftReduceField()
LABELS = data.Field(sequential=False)
train, dev, test = datasets.SNLI.splits(
TEXT, TRANSITIONS, LABELS, wv_type='glove.42B')
TEXT.build_vocab(train, dev, test)
train_iter, dev_iter, test_iter = data.BucketIterator.splits(
(train, dev, test), batch_size=64)

你可以在 train.py 找到其余代码，包括训练循环（loop）的和衡量精度的。现在讲模型。如同上篇所描述，一个 SPINN 编码器包含一个参数化的 Reduce 层，以及可选的 recurrent Tracker，以追踪语境。这通过在神经网络每读取一个词语、或应用 Reduce 的时候，更新隐藏状态来实现。下面的代码其实表示了，创建一个 SPINN 只是意味着创建这两个子模块而已，以及把它们放到容器里面以日后使用。

import torch
from torch import nn
# subclass the Module class from PyTorch’s neural network package
class SPINN(nn.Module):
def __init__(self, config):
super(SPINN, self).__init__()
self.config = config
self.reduce = Reduce(config.d_hidden, config.d_tracker)
if config.d_tracker is not None:
self.tracker = Tracker(config.d_hidden, config.d_tracker)

创建模型时，SPINN.__init__被调用一次。它分配、初始化参数，但不进行任何神经网络运算，也不涉及创建计算图。每组新数据 batch 上运行的代码，在 SPINN 中定义。PyTorch 里，用户定义模型前馈通道的方法名为 “forward”。事实上，它是对上文提到的 stack-manipulation 算法的实现，在普通 Python 里，它运行于 Buffer 和堆栈的 batch 上——对每个样例使用两者之一。 在转换过程包含的“shift” 和 “reduce” op 上迭代，如果它存在，就运行 Tracker，并运行于 batch 中的每个样例以应用 “shift”op，或加入需要 “reduce” op 的样例列表。然后在列表所有的样例上运行 Reduce 层，把结果 push 回相关堆栈。

def forward(self, buffers, transitions):
# The input comes in as a single tensor of word embeddings;
# I need it to be a list of stacks, one for each example in
# the batch, that we can pop from independently. The words in
# each example have already been reversed, so that they can
# be read from left to right by popping from the end of each
# list; they have also been prefixed with a null value.
buffers = [list(torch.split(b.squeeze(1), 1, 0))
for b in torch.split(buffers, 1, 1)]
# we also need two null values at the bottom of each stack,
# so we can copy from the nulls in the input; these nulls
# are all needed so that the tracker can run even if the
# buffer or stack is empty
stacks = [[buf[0], buf[0]] for buf in buffers]
if hasattr(self, 'tracker'):
self.tracker.reset_state()
for trans_batch in transitions:
if hasattr(self, 'tracker'):
# I described the Tracker earlier as taking 4
# arguments (context_t, b, s1, s2), but here I
# provide the stack contents as a single argument
# while storing the context inside the Tracker
# object itself.
tracker_states, _ = self.tracker(buffers, stacks)
else:
tracker_states = itertools.repeat(None)
lefts, rights, trackings = [], [], []
batch = zip(trans_batch, buffers, stacks, tracker_states)
for transition, buf, stack, tracking in batch:
if transition == SHIFT:
stack.append(buf.pop())
elif transition == REDUCE:
rights.append(stack.pop())
lefts.append(stack.pop())
trackings.append(tracking)
if rights:
reduced = iter(self.reduce(lefts, rights, trackings))
for transition, stack in zip(trans_batch, stacks):
if transition == REDUCE:
stack.append(next(reduced))
return [stack.pop() for stack in stacks]

调用 self.tracker 或 self.reduce，会相对应地运行 Tracker 中的“forward”方式，或 Reduce 子模块。这需要在一个样例列表来执行该 op。所有数学运算密集、用 GPU 加速、收益用 batch 的 op 都发生在  Tracker 和 Reduce 之中。因此，在主要的“forward”方式中，单独在不同样例上运行；对 batch 中的每个样例保持独立的 buffer 和堆栈，都是意义的。为了更干净地写这些函数，我会用一些辅助，把这些样例列表转为 batch 化的张量，反之亦然。

我倾向于让 Reduce 模块自动 batch 参数来加速计算，然后 unbatch 它们，这样之后能单独地 push、pop。把每一组左右子短语放到一起，来表示母短语的合成函数是 TreeLSTM，一个常规 LSTM 的变种。此合成函数要求，所有子树的状态要由两个张量组成，一个隐藏状态 h 和一个内存单元状态 c。定义该函数的因素有两个：运行于子树隐藏状态中的两个线性层  (nn.Linear)，以及非线性合成函数 tree_lstm，后者把线性层的结果和子树内存单元的状态组合起来。在 SPINN 中，这通过加入第三个运行于 Tracker 隐藏状态的 线性层来拓展。

def tree_lstm(c1, c2, lstm_in):
# Takes the memory cell states (c1, c2) of the two children, as
# well as the sum of linear transformations of the children’s
# hidden states (lstm_in)
# That sum of transformed hidden states is broken up into a
# candidate output a and four gates (i, f1, f2, and o).
a, i, f1, f2, o = lstm_in.chunk(5, 1)
c = a.tanh() * i.sigmoid() + f1.sigmoid() * c1 + f2.sigmoid() * c2
h = o.sigmoid() * c.tanh()
return h, c

class Reduce(nn.Module):
def __init__(self, size, tracker_size=None):
super(Reduce, self).__init__()
self.left = nn.Linear(size, 5 * size)
self.right = nn.Linear(size, 5 * size, bias=False)
if tracker_size is not None:
self.track = nn.Linear(tracker_size, 5 * size, bias=False)

def forward(self, left_in, right_in, tracking=None):
left, right = batch(left_in), batch(right_in)
tracking = batch(tracking)
lstm_in = self.left(left[0])
lstm_in += self.right(right[0])
if hasattr(self, 'track'):
lstm_in += self.track(tracking[0])
return unbatch(tree_lstm(left[1], right[1], lstm_in))

由于 Reduce 层和以与之类似方式执行的 Tracker 都在 LSTM 上运行，batch 和 unbatch 辅助函数会在成对隐藏、内存状态上运行。

def batch(states):
if states is None:
return None
states = tuple(states)
if states[0] is None:
return None
# states is a list of B tensors of dimension (1, 2H)
# this returns two tensors of dimension (B, H)
return torch.cat(states, 0).chunk(2, 1)

def unbatch(state):
if state is None:
return itertools.repeat(None)
# state is a pair of tensors of dimension (B, H)
# this returns a list of B tensors of dimension (1, 2H)
return torch.split(torch.cat(state, 1), 1, 0)