LSTM JIT

本文探讨了Pytorch模型的C++部署过程,重点介绍了如何利用torch.jit进行模型封装,以便直接在C++环境中调用。文章分享了调试技巧、lstm模块的处理方法及模型保存和加载的实现。
部署运行你感兴趣的模型镜像

前言

项目部署考虑到加密,算法是基于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.youkuaiyun.com/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")

您可能感兴趣的与本文相关的镜像

PyTorch 2.5

PyTorch 2.5

PyTorch
Cuda

PyTorch 是一个开源的 Python 机器学习库,基于 Torch 库,底层由 C++ 实现,应用于人工智能领域,如计算机视觉和自然语言处理

### 如何部署 LSTM 模型 #### PyTorch 下从头搭建并部署 LSTM 模型 在 PyTorch 中,LSTM 的实现可以通过自定义模型类完成。以下是基于引用的内容以及专业知识总结的一个完整的解决方案。 首先,在 PyTorch 中创建一个基本的 LSTM 模型: ```python import torch import torch.nn as nn class LSTMModel(nn.Module): def __init__(self, input_size, hidden_size, num_layers, output_size): super(LSTMModel, self).__init__() self.hidden_size = hidden_size self.num_layers = num_layers # 定义 LSTM 层 self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True) # 定义全连接层用于输出 self.fc = nn.Linear(hidden_size, output_size) def forward(self, x): h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) out, _ = self.lstm(x, (h0, c0)) out = self.fc(out[:, -1, :]) # 只取最后一个时间步的结果 return out ``` 此代码片段展示了一个基础的 LSTM 构建方法[^1]。通过 `nn.LSTM` 实现多层 LSTM 结构,并利用线性层将隐藏状态映射到最终预测值。 --- #### 部署 LSTM 模型的方法 当训练完成后,可以采用多种方式来部署该模型。以下是一些常见的部署方案及其具体操作说明: ##### 使用 TorchScript 进行序列化 PyTorch 提供了一种称为 **TorchScript** 的机制,允许开发者将 Python 脚本转换为独立于解释器的形式以便高效运行。这一步骤对于 C++ 或其他环境中的推理至关重要。 保存模型为 TorchScript 文件: ```python model_scripted = torch.jit.script(model) # 将模型导出为 TorchScript 格式 model_scripted.save('lstm_model.pt') # 保存文件 ``` 加载已保存的模型(适用于生产环境中): ```cpp #include <torch/script.h> std::shared_ptr<torch::jit::script::Module> module = torch::jit::load("lstm_model.pt"); ``` 上述过程展示了如何使用 PyTorch 自带工具链将模型迁移到 C++ 平台[^2]。 ##### 导出 ONNX 格式的模型 另一种流行的方式是将模型导出为 Open Neural Network Exchange (**ONNX**) 格式,这是一种跨框架的标准表示形式。它支持广泛的硬件加速库和推理引擎。 导出流程如下所示: ```python dummy_input = torch.randn(1, seq_length, input_dim) # 创建虚拟输入张量 torch.onnx.export( model, dummy_input, "lstm_model.onnx", export_params=True, opset_version=11, do_constant_folding=True, input_names=['input'], output_names=['output'] ) ``` 之后可以在诸如 TensorRT、OpenVINO 等高性能计算平台上进一步优化性能[^3]。 --- #### 性能调优建议 为了提高部署后的效率,还可以考虑以下几个方面: - 对权重矩阵应用量化技术减少存储需求; - 利用 GPU/CPU 多核特性提升批量处理能力; - 如果目标平台资源有限,则裁剪不必要的神经元节点简化架构设计。 ---
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

勤劳的凌菲

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值