高并行递归的简单递归单元
sru的Python项目详细描述
关于
sru是一个循环单元,其运行速度比cudnn lstm快10倍以上,在许多任务中测试的准确性没有损失。
在gtx 1070上测试的lstm、conv2d和sru的平均处理时间
paper有多个版本,请检查最新版本。
参考:
高度并行递归的简单递归单元
@inproceedings{lei2018sru,
title={Simple Recurrent Units for Highly Parallelizable Recurrence},
author={Tao Lei and Yu Zhang and Sida I. Wang and Hui Dai and Yoav Artzi},
booktitle={Empirical Methods in Natural Language Processing (EMNLP)},
year={2018}
}
要求
- PyTorch>;=0.4.1推荐,pytorch installation details
- CuPy
- pynvrtc
- ninja(可选)用于CPU上的快速推理。
通过pip install -r requirements.txt
安装需求。Cupy和Pynvrtc需要支持GPU上的培训/测试。
安装
来源:
sru可以通过python setup.py install
或pip install .
作为常规包安装。
来自PYPI:
pip install sru
pip install sru[cuda]
另外安装cupy和pynvrc。
pip install sru[cpu]
另外安装忍者
直接使用源而不安装:
确保系统可以找到此repo和cuda库,例如
export PYTHONPATH=path_to_repo/sru
export LD_LIBRARY_PATH=/usr/local/cuda/lib64
示例
sru的用法类似于nn.LSTM
。SRU可能需要比LSTM更多的堆叠层。我们建议从2层开始,必要时使用更多(更多实验细节见我们的报告)。
importtorchfromtorch.autogradimportVariablefromsruimportSRU,SRUCell# input has length 20, batch size 32 and dimension 128x=Variable(torch.FloatTensor(20,32,128).cuda())input_size,hidden_size=128,128rnn=SRU(input_size,hidden_size,num_layers=2,# number of stacking RNN layersdropout=0.0,# dropout applied between RNN layersbidirectional=False,# bidirectional RNNlayer_norm=False,# apply layer normalization on the output of each layerhighway_bias=0,# initial bias of highway gate (<= 0)rescale=True,# whether to use scaling correction)rnn.cuda()output_states,c_states=rnn(x)# forward pass# output_states is (length, batch size, number of directions * hidden size)# c_states is (layers, batch size, number of directions * hidden size)
贡献者
https://github.com/taolei87/sru/graphs/contributors
其他实现
@musyoku有一个非常漂亮的SRU implementaion链子。
@adrianbg实现了第一个CPU version。