Pytorch 1.1.0，來了！

盼望已久，Pytorch終于更新了！Pytroch 1.1.0的發(fā)布除了修復(fù)了已有bug之外，最大的亮點(diǎn)就是可以更快、更好的支持自定義RNN，以及TensorBoard對(duì)可視化和模型調(diào)試提供了一流的本地支持。

Pytorch 1.1.0，來了！

可以說是一大波更新來襲了，話不多說上亮點(diǎn)：

TorchScript(Pytorch JIT)更快、更好的支持自定義RNN；

TensorBoard對(duì)可視化和模型調(diào)試提供了一流的本地支持；

可以在ScriptModule上通過使用torch.jit包裝屬性來分配屬性；

TorchScript現(xiàn)在對(duì)列表和字典類型提供了魯棒性的支持；

對(duì)于更復(fù)雜的有狀態(tài)操作，TorchScript現(xiàn)在支持使用@torch.jit.script注釋類；

nn.parallel.DistributedDataParallel：現(xiàn)在可以包裝多GPU模塊，它可以在一臺(tái)服務(wù)器上實(shí)現(xiàn)模型并行和跨服務(wù)器的數(shù)據(jù)并行等用例。

注：不再支持CUDA 8.0。

此更新一出，在Reddit上也引發(fā)了一波熱議，大部分網(wǎng)友們表示：

“贊！”、“好用！”、“愛了！”

用TorchScript優(yōu)化CUDA遞歸神經(jīng)網(wǎng)絡(luò)

Pytorch添加的一個(gè)新特性是更好地支持帶有TorchScript (PyTorch JIT)的快速自定義遞歸神經(jīng)網(wǎng)絡(luò)(fastrnns)。

RNN是一種流行的模型，在各種NLP任務(wù)上都表現(xiàn)出了良好的性能。PyTorch可以實(shí)現(xiàn)許多最流行的變體，例如Elman RNN、GRU和LSTM，以及多層和雙向變體。

然而，許多用戶希望實(shí)現(xiàn)他們自己的自定義RNN。將層規(guī)范化應(yīng)用于LSTM就是這樣一種用例。由于PyTorch CUDA LSTM實(shí)現(xiàn)使用融合內(nèi)核，因此很難插入規(guī)范化甚至修改基本LSTM實(shí)現(xiàn)。

許多用戶已經(jīng)轉(zhuǎn)向使用標(biāo)準(zhǔn)PyTorch運(yùn)算符編寫自定義實(shí)現(xiàn)，但是這樣的代碼遭受高開銷：大多數(shù)PyTorch操作在GPU上啟動(dòng)至少一個(gè)內(nèi)核，并且RNN由于其重復(fù)性質(zhì)通常運(yùn)行許多操作。但是可以應(yīng)用TorchScript來融合操作并自動(dòng)優(yōu)化代碼，在GPU上啟動(dòng)更少、更優(yōu)化的內(nèi)核。

此次更新的目標(biāo)之一是讓用戶能夠在TorchScript中編寫快速，自定義的RNN，而無需編寫專門的CUDA內(nèi)核來實(shí)現(xiàn)類似的性能。接下來將提供如何使用TorchScript編寫自己的快速RNN的教程。

編寫自定義RNN

首先，可以使用下方鏈接中的文件作為模板來編寫自己的自定義RNN。

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

如果想獲得TorchScript當(dāng)前提供的速度/優(yōu)化（如運(yùn)算符融合，批量矩陣乘法等），請(qǐng)遵循以下指南。

如果定制操作都是element-wise的，那就可以自動(dòng)獲得PyTorch JIT操作符fusion的優(yōu)勢！

如果有更復(fù)雜的操作（例如，reduce和element-wise的渾南和操作），請(qǐng)考慮分別對(duì)reduce操作和element-wise操作進(jìn)行分組。

如果想知道自定義RNN中融合了什么，可以使用graph_for檢查操作的優(yōu)化圖。以LSTMCell為例：

#getinputsandstatesforLSTMCellinputs=get_lstm_inputs()#instantiateaScriptModulecell=LSTMCell(input_size,hidden_size)#printtheoptimizedgraphusinggraph_forout=cell(inputs)print(cell.graph_for(inputs))

這將提供的專用輸入生成優(yōu)化的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)

從上圖中可以看到它有一個(gè)prim :: FusionGroup_0子圖，它融合了LSTMCell中的所有element-wise操作（轉(zhuǎn)置和矩陣乘法不是element-wise操作）。

可變長度序列最佳實(shí)踐

TorchScript不支持PackedSequence。 通常，當(dāng)處理可變長度序列時(shí)，最好將它們填充到單個(gè)張量中并通過TorchScript LSTM發(fā)送該張量。 例如：

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

當(dāng)然，output可能在填充區(qū)域中有一些垃圾數(shù)據(jù)；使用lengths來跟蹤你不需要的部分。

優(yōu)化

現(xiàn)在將解釋PyTorch JIT為加速自定義RNN所執(zhí)行的優(yōu)化。 將在TorchScript中使用一個(gè)簡單的自定義LSTM模型來說明優(yōu)化，但其中許多是通用的并適用于其他RNN。

為了說明所做的優(yōu)化以及如何從這些優(yōu)化中獲益，將運(yùn)行一個(gè)用TorchScript編寫的簡單自定義LSTM模型（可以參考custom_lstm.py中的代碼或下面的代碼片段）并計(jì)算更改。

在配備2個(gè)Intel Xeon芯片和一個(gè)Nvidia P100的機(jī)器中設(shè)置環(huán)境，安裝了cuDNN v7.3，CUDA 9.2。 LSTM模型的基本設(shè)置如下：

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

PyTorch JIT最重要的是將python程序編譯為PyTorch JIT IR，這是一個(gè)用于對(duì)程序圖形結(jié)構(gòu)進(jìn)行建模的中間表示。然后，該IR可以從整個(gè)程序優(yōu)化，硬件加速中受益，并且總體上具有提供大量計(jì)算增益的潛力。

接下來，將解釋在如何提高訓(xùn)練或推理性能方面所做的主要優(yōu)化，從LSTMCell和LSTMLayer開始，以及一些misc優(yōu)化。

LSTM Cell(前向)

LSTM中的幾乎所有計(jì)算都發(fā)生在LSTMCell中，因此重要的是看看它包含的計(jì)算以及如何提高它們的速度。 下面是TorchScript中的LSTMCell實(shí)現(xiàn)示例：

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）可實(shí)現(xiàn)多種優(yōu)化和可伸縮計(jì)算。 除了可以做的典型編譯器優(yōu)化（CSE，常量傳播等）之外，還可以運(yùn)行其他IR轉(zhuǎn)換以使代碼運(yùn)行得更快。

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上做了一些技巧來提高性能，團(tuán)隊(duì)做了一些示例優(yōu)化：

循環(huán)展開(Loop Unrolling)：自動(dòng)在代碼中展開循環(huán)（對(duì)于大循環(huán)，展開它的一小部分），然后授權(quán)對(duì)for循環(huán)控制流進(jìn)行進(jìn)一步的優(yōu)化。 例如，fuser可以將循環(huán)體的迭代中的操作融合在一起，這導(dǎo)致對(duì)于諸如LSTM的控制流密集型模型的良好性能改進(jìn)。

批量矩陣乘法：對(duì)于輸入預(yù)乘的RNN（即模型具有大量相同LHS或RHS的矩陣乘法），可以將這些操作一起有效地批量處理為單個(gè)矩陣乘法，同時(shí)對(duì)輸出進(jìn)行分塊以實(shí)現(xiàn)等效語義。

通過應(yīng)用這些技術(shù)，將前向傳播的時(shí)間減少了1.6ms，達(dá)到8.4ms(1.2倍加速)，后向傳播的時(shí)間減少了7ms，達(dá)到20ms左右(1.35倍加速)。

LSTM層(后向)

“樹結(jié)構(gòu)”批處理矩陣Muplication：通常情況是在LSTM反向圖中多次重復(fù)使用單個(gè)權(quán)重，形成一個(gè)樹，其中葉子是矩陣乘法，節(jié)點(diǎn)是相加的。 這些節(jié)點(diǎn)可以通過在不同維度上連接LHS和RHS來組合在一起，然后計(jì)算為單個(gè)矩陣乘法。 等價(jià)公式可表示如下：

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

Autograd是使PyTorch成為如此優(yōu)雅的ML框架的關(guān)鍵組件。因此，將其應(yīng)用到PyTorch JIT，但是使用了一種新的自動(dòng)微分(AD)機(jī)制，該機(jī)制在IR級(jí)別上工作。JIT自動(dòng)微分將把正向圖分割成符號(hào)可微分的子圖，并為這些子圖生成向后節(jié)點(diǎn)。以上面的IR為例，對(duì)于具有AD公式的操作，我們將圖節(jié)點(diǎn)分組為一個(gè)prim :: DifferentiableGraph_0。對(duì)于沒有添加到AD公式中的操作，我們將在執(zhí)行期間返回到Autograd。

優(yōu)化反向路徑是困難的，隱式broadcasting語義使得自動(dòng)微分的優(yōu)化更加困難。 PyTorch可以方便地編寫張量操作，而無需通過broadcasting張量來擔(dān)心形狀。 對(duì)于性能而言，反向的痛點(diǎn)是需要對(duì)這種可broadcasting操作進(jìn)行求和。 這導(dǎo)致每個(gè)可broadcasting操作的導(dǎo)數(shù)后跟一個(gè)求和。 由于目前無法融合減少操作，這會(huì)導(dǎo)致FusionGroups分成多個(gè)小組，從而導(dǎo)致性能下降。 要解決這個(gè)問題，請(qǐng)參閱Thomas Viehmann撰寫的文章：http://lernapparat.de/fast-lstm-pytorch/。

更多這方面的優(yōu)化內(nèi)容可參考Pytorch團(tuán)隊(duì)博客原文：

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

更多新功能

運(yùn)算符

torch.tril_indices, torch.triu_indices:添加了與NumPy具有相同行為的運(yùn)算符；

torch.combinations, torch.cartesian_prod:添加了類似于itertools的新運(yùn)算符；

torch.repeat_interleave:新運(yùn)算符類似于numpy.repeat；

torch.from_file:類似于Storage.from_file的新運(yùn)算符，但返回一個(gè)張量；

torch.unique_consecutive:新的運(yùn)算符，其語義類似于C ++中的std :: unique；

torch.tril, torch.triu, torch.trtrs:現(xiàn)在支持批處理；

torch.gather:添加對(duì)sparse_grad選項(xiàng)的支持；

torch.std, torch.max_values, torch.min_values, torch.logsumexp現(xiàn)在可以同時(shí)在多個(gè)維度上運(yùn)行；

torch.cdist:添加了與scipy.spatial.distance.cdist等效的運(yùn)算符；

torch.__config__.show():報(bào)告所有庫的詳細(xì)版本。

NN

nn.MultiheadedAttention:從注意力中實(shí)現(xiàn)MultiheadedAttention的新模塊；

nn.functional.interpolate:增加了對(duì)bicubic的支持；

nn.SyncBatchNorm:支持同步批量標(biāo)準(zhǔn)化；

nn.Conv:通過mode ='circular'添加了對(duì)Circular Padding的支持；

nn.EmbeddingBag:現(xiàn)在支持可訓(xùn)練的`per_sample_weights；

nn.EmbeddingBag:添加對(duì)from_pretrained方法的支持，如nn.Embedding中所示；

RNNs:通過enforce_sorted自動(dòng)處理未排序的可變長度序列；

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

更多有關(guān)張量/dtypes、性能提高、bug修復(fù)、棄用的項(xiàng)目等內(nèi)容可查看Pytorch在GitHub發(fā)布的項(xiàng)目原文：

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