Pytorch C++部署 之 TorchScript 踩坑记录

小白一枚,最近尝试用c++调pytorch模型,网上搜了一对资料,大同小异,都是需要先将模型文件转换成torch script。
pytorch官网提供两种方式,1:trace 2:注释。
这两种方式怎么操作,官网以及很多博客都有详细介绍,不细扣,我也扣不明白,毕竟我是个小白,放链接供参考:
https://blog.csdn.net/dou3516/article/details/82912480

============================
以下记录踩坑记录:
个人环境:
win10
cuda9.0 + cudnn7.0.5
1050Ti-4G
ptorch 1.0稳定版。。。
默认各种环境已经装好,开始撸起袖子码代码。

自己pytorch写了个简单的分类网络,通过注释的方式将模型序列化,举个栗子:

class myModel(torch.jit.ScriptModule):
    def __init__(self,num_classes):
            super(myModel, self).__init__() 
            ...........         
    @torch.jit.script_method
    def forward(self, x):
             ..........
          return x
 my_script_module = myModel(2)
 my_script_module.save("model.pt")

反正按照官方说法,加注释就是按这个样子加的,那么问题来了,各种各样报错。


错误1:ValueError: substring not found
原因:forward 函数里不能有中文注释


错误2: torch.jit.frontend.Unsuppo

<think>嗯,用户想了解如何使用TorchScript进行PyTorch模型的部署。首先,我需要回忆一下TorchScript的基本概念和用途。根据引用[4]的内容,TorchScript是Python的一个子集,主要用于将PyTorch代码转换成C++代码,提高部署效率。那用户可能关心的是具体转换步骤和注意事项。 接下来,用户可能需要知道两种主要的转换方法:torch.jit.trace和torch.jit.script。根据引用[3]和引用[5],trace适用于没有控制流的模型,而script可以处理更复杂的逻辑。需要详细说明这两种方法的区别和使用场景。 然后,用户可能会遇到实际应用中的问题,比如如何处理动态控制流。这时候需要引用例子,比如引用[5]中的代码示例,展示如何用@torch.jit.script_method装饰器来定义forward方法。同时,要提醒用户注意输入输出的类型一致性,避免转换失败。 另外,模型保存和加载也是关键步骤。需要提到使用save方法保存为.pt文件,以及如何在C++中加载,如引用[4]提到的ScriptModule格式。这部分可能需要分步骤说明,确保用户能按照流程操作。 用户可能还会关心性能优化和实际应用中的好处。引用[1]提到TorchScript可以保护模型知识产权,这点需要强调,同时引用[2]推荐的相关资料可能对用户深入学习有帮助。 最后,生成相关问题时要覆盖常见疑问,比如trace和script的区别、动态控制流的处理、性能优化方法以及实际应用场景。这样用户可以根据自己的需求进一步探索。</think>### 如何使用 TorchScript 进行 PyTorch 模型部署 #### 1. TorchScript 的核心作用 - **跨平台部署**:将动态图模型转换为静态图(ScriptModule),支持脱离 Python 环境运行(如 C++/Android)[^4]。 - **性能优化**:通过 JIT 编译提升推理效率,减少 Python 解释器开销[^1]。 - **模型保护**:固化模型结构和参数,避免源码泄露[^1]。 #### 2. 模型转换方法 ##### 方法一:`torch.jit.trace`(适用于无控制流的模型) ```python import torch # 定义模型 class SimpleModel(torch.nn.Module): def forward(self, x): return x * 2 model = SimpleModel() # 输入示例 example_input = torch.rand(1, 3) # 跟踪执行流程 traced_model = torch.jit.trace(model, example_input) traced_model.save("model_traced.pt") ``` *特点*:仅记录特定输入下的计算路径,无法处理动态控制流[^3]。 ##### 方法二:`torch.jit.script`(支持动态逻辑) ```python class DynamicModel(torch.nn.Module): def forward(self, x): if x.sum() > 0: return x * 3 else: return x * (-1) scripted_model = torch.jit.script(DynamicModel()) scripted_model.save("model_scripted.pt") ``` *特点*:直接解析 Python 代码,支持条件分支/循环等动态结构[^3]。 #### 3. 混合使用技巧 - 对包含复杂逻辑的子模块使用`@torch.jit.script`装饰器 - 对数据流明确的模块使用`torch.jit.trace` ```python class HybridModel(torch.nn.Module): def __init__(self): super().__init__() self.linear = torch.nn.Linear(10, 5) @torch.jit.script_method def process(self, x): # 包含复杂逻辑 return torch.relu(x) def forward(self, x): x = self.linear(x) return self.process(x) model = torch.jit.script(HybridModel()) ``` #### 4. 部署验证要点 1. **输入类型约束**:确保部署时输入张量类型/维度与转换时一致 2. **跨设备兼容性**:转换前使用`model.to("cpu")`避免 GPU 依赖 3. **版本一致性**:保证 PyTorch 版本在转换/部署环境中的一致性 #### 5. C++ 加载示例 ```cpp #include <torch/script.h> int main() { torch::jit::script::Module module; module = torch::jit::load("model.pt"); std::vector<torch::jit::IValue> inputs; inputs.push_back(torch::ones({1, 3})); at::Tensor output = module.forward(inputs).toTensor(); } ```
评论 13
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值