pytorch模型转成java_PyTorch 1.0 中文官方教程：在C++中加载PYTORCH模型

在C++中加载PYTORCH模型

PyTorch的主要接口为Python。虽然Python有动态编程和易于迭代的优势，但在很多情况下，正是Python的这些属性会带来不利。我们经常遇到的生产环境，要满足低延迟和严格部署要求。对于生产场景而言，C++通常是首选语言，也能很方便的将其绑定到另一种语言，如Java，Rust或Go。本教程将介绍从将PyTorch训练的模型序列化表示，到C++语言_加载_和_执行_的过程。

第一步：将PyTorch模型转换为Torch Script

PyTorch模型从Python到C++的转换由Torch Script实现。Torch Script是PyTorch模型的一种表示，可由Torch Script编译器理解，编译和序列化。如果使用基础的“eager”API编写的PyTorch模型，则必须先将模型转换为Torch Script，当然这也是比较容易的。如果已有模型的Torch Script，则可以跳到本教程的下一部分。

将PyTorch模型转换为Torch Script有两种方法。

第一种方法是Tracing。该方法通过将样本输入到模型中一次来对该过程进行评估从而捕获模型结构.并记录该样本在模型中的flow。该方法适用于模型中很少使用控制flow的模型。

第二个方法就是向模型添加显式注释(Annotation)，通知Torch Script编译器它可以直接解析和编译模型代码，受Torch Script语言强加的约束。

小贴士

可以在官方的Torch Script 参考中找到这两种方法的完整文档，以及有关使用哪个方法的细节指导。

利用Tracing将模型转换为Torch Script

要通过tracing来将PyTorch模型转换为Torch脚本,必须将模型的实例以及样本输入传递给torch.jit.trace函数。这将生成一个 torch.jit.ScriptModule对象，并在模块的forward方法中嵌入模型评估的跟踪：

import torch

import torchvision

# 获取模型实例

model = torchvision.models.resnet18()

# 生成一个样本供网络前向传播 forward()

example = torch.rand(1, 3, 224, 224)

# 使用 torch.jit.trace 生成 torch.jit.ScriptModule 来跟踪

traced_script_module = torch.jit.trace(model, example)

现在，跟踪的ScriptModule可以与常规PyTorch模块进行相同的计算：

In[1]: output = traced_script_module(torch.ones(1, 3, 224, 224))

In[2]: output[0, :5]

Out[2]: tensor([-0.2698, -0.0381, 0.4023, -0.3010, -0.0448], grad_fn=)

通过Annotation将Model转换为Torch Script

在某些情况下，例如，如果模型使用特定形式的控制流，如果想要直接在Torch Script中编写模型并相应地标注(annotate)模型。例如，假设有以下普通的 Pytorch模型：

import torch

class MyModule(torch.nn.Module):

def __init__(self, N, M):

s