LSTM JIT

前言

项目部署考虑到加密，算法是基于Pytorch进行开发的，所以决定利用Pytorch的新功能torch.jit进行封装，以便直接将模型用C++调用。

参考

Pytorch官方torch.git教程：https://pytorch.org/docs/stable/jit.html
Pytorch官方基于torch.jit 的Seq2Seq示例：https://pytorch.org/tutorials/beginner/deploy_seq2seq_hybrid_frontend_tutorial.html
还有大神翻译的https://blog.csdn.net/xxradon/article/details/88063320
官方jit的测试代码https://github.com/pytorch/pytorch/blob/master/test/test_jit.py

调试那点小事儿

改写代码后，调试代码应该设置环境变量PYTORCH_JIT=0，pycharm目前不知道怎么设置环境变量，所采用的方法是将torch.jit中的__init__.py打开，修改以下代码

_enabled = _parse_env('PYTORCH_JIT', True, "> Using PyTorch JIT", "> PyTorch JIT DISABLED")

将True改为False则对应PYTORCH_JIT=0，可以方便的进行调试或者关闭jit环境。

lstm模块

按照我的理解，trace模式时，会将example_inputs执行一遍，然后将整个状态记录下来，也就意味着lstm只能是一定的长度，而我们的问答系统输入问题的长度是可变的，需要可变长度的lstm，所以应该不适用于trace模式，因此我将lstm所在的模型改为script模式，结果就报各种错误，查看基于torch.jit 的Seq2Seq示例和官方jit的测试代码，发现其中的lstm都是采用trace模式来做的，对于lstm模块的代码未进行任何改动，在使用时

 # 加载model
encoder_model = net.Encoder(self.numLayers, embeddingSize=self.embeddingSize,
                                     hiddenSize=self.hiddenSize, word2id=self.word2id)
        
encoder_model.to(device)
encoder_model.eval()
encoder_model.load_state_dict(torch.load(model_dir))
# Create artificial inputs
#以下两个tensor是模拟forward所需要的两个输入
test_seq = torch.LongTensor(1, 5).random_(0, len(self.word2id)).to(device) 
test_seq_length = torch.LongTensor([test_seq.size()[0]]).to(device)
trace_encoder_model = torch.jit.trace(encoder_model, (test_seq, test_seq_length))
trace_encoder_model.save("model.pt")

Encoder即为要进行加密的网络，这样就将我们的网络保存为一个model.pt文件了，

使用

– python接口

encoder_model = torch.jit.load("model.pt")

–c++接口

encoder_model = torch::jit::load("model.pt")