• 1
  • 2
  • 3
  • 4

首页 / 行业

Pytorch 1.1.0,来了!

2019-05-05 10:02:00

盼望已久,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为例:

#getinputsandstatesforLSTMCellinputs=get_lstm_inputs()#instantiateaScriptModulecell=LSTMCell(input_size,hidden_size)#printtheoptimizedgraphusinggraph_forout=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)withprim::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)...somebroadcastsizesoperations...%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)...somebroadcastsizesoperations...return(%hy,%cy,%49,%44,%196,%199,%340,%192,%325,%185,%ingate.1,%forgetgate.1,%cellgate.1,%outgate.1,%395,%396,%287)withprim::FusionGroup_0=graph(%13:Float(*,*),%71:Tensor,%76:Tensor,%81:Tensor,%86:Tensor):...somechunks,constants,andaddoperations...%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],eachTensorisT'xCpadded=torch.utils.rnn.pad_sequence(sequences)lengths=[seq.size(0)forseqinsequences]padded#TxNxC,whereNisbatchsizeandTisthemaxofallT'model=LSTM(...)output,hiddens=model(padded)output#TxNxC

当然,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=512hidden_size=512mini_batch=64numLayers=1seq_length=100

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

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

LSTM Cell(前向)

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

classLSTMCell(jit.ScriptModule):def__init__(self,input_size,hidden_size):super(LSTMCell,self).__init__()self.input_size=input_sizeself.hidden_size=hidden_sizeself.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_methoddefforward(self,input,state):#type:(Tensor,Tuple[Tensor,Tensor])->Tuple[Tensor,Tuple[Tensor,Tensor]]hx,cx=stategates=(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)returnhy,(hy,cy)

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

LSTM层(前向)

classLSTMLayer(jit.ScriptModule):def__init__(self,cell,*cell_args):super(LSTMLayer,self).__init__()self.cell=cell(*cell_args)@jit.script_methoddefforward(self,input,state):#type:(Tensor,Tuple[Tensor,Tensor])->Tuple[Tensor,Tuple[Tensor,Tensor]]inputs=input.unbind(0)outputs=torch.jit.annotate(List[Tensor],[])foriinrange(len(inputs)):out,state=self.cell(inputs[i],state)outputs+=[out]returntorch.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/

更多新功能

运算符

torch.tril_indices, torch.triu_indices:添加了与NumPy具有相同行为的运算符;

torch.combinations, torch.cartesian_prod:添加了类似于itertools的新运算符;

torch.repeat_interleave:新运算符类似于numpy.repeat;

torch.from_file:类似于Storage.from_file的新运算符,但返回一个张量;

torch.unique_consecutive:新的运算符,其语义类似于C ++中的std :: unique;

torch.tril, torch.triu, torch.trtrs:现在支持批处理;

torch.gather:添加对sparse_grad选项的支持;

torch.std, torch.max_values, torch.min_values, torch.logsumexp现在可以同时在多个维度上运行;

torch.cdist:添加了与scipy.spatial.distance.cdist等效的运算符;

torch.__config__.show():报告所有库的详细版本。

NN

nn.MultiheadedAttention:从注意力中实现MultiheadedAttention的新模块;

nn.functional.interpolate:增加了对bicubic的支持;

nn.SyncBatchNorm:支持同步批量标准化;

nn.Conv:通过mode ='circular'添加了对Circular Padding的支持;

nn.EmbeddingBag:现在支持可训练的`per_sample_weights;

nn.EmbeddingBag:添加对from_pretrained方法的支持,如nn.Embedding中所示;

RNNs:通过enforce_sorted自动处理未排序的可变长度序列;

nn.Identity:便于模型surgery的新模块。

更多有关张量/dtypes、性能提高、bug修复、弃用的项目等内容可查看Pytorch在GitHub发布的项目原文:

https://github.com/pytorch/pytorch/releases/tag/v1.1.0

启动内核操作开销

  • 1
  • 2
  • 3
  • 4

最新内容

手机

相关内容

  • 1
  • 2
  • 3

猜你喜欢