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")

猜你喜欢

转载自blog.csdn.net/Suan2014/article/details/88234234
JIT