一、模型保存与加载

三个核心函数：

1. torch.save ：把序列化的对象保存到硬盘。它利用了 Python 的 pickle 来实现序列化。模型、张量以及字典都可以用该函数进行保存；
2. torch.load：采用 pickle 将反序列化的对象从存储中加载进来。

3. torch.nn.Module.load_state_dict：采用一个反序列化的 state_dict加载一个模型的参数字典。

   1. 什么是状态字典(state_dict)

   PyTorch 中，一个模型(torch.nn.Module)的可学习参数(也就是权重和偏置值)是包含在模型参数(model.parameters())中的，一个状态字典就是一个简单的 Python 的字典，其键值对是每个网络层和其对应的参数张量。模型的状态字典只包含带有可学习参数的网络层（比如卷积层、全连接层等）和注册的缓存（batchnorm的 running_mean）。优化器对象(torch.optim)同样也是有一个状态字典，包含的优化器状态的信息以及使用的超参数。
   由于状态字典也是 Python 的字典，因此对 PyTorch 模型和优化器的保存、更新、替换、恢复等操作都很容易实现。
   下面是一个简单的使用例子 ```python

   Define model

   class TheModelClass(nn.Module): def init(self):

    super(TheModelClass, self).__init__()
    self.conv1 = nn.Conv2d(3, 6, 5)
    self.pool = nn.MaxPool2d(2, 2)
    self.conv2 = nn.Conv2d(6, 16, 5)
    self.fc1 = nn.Linear(16 * 5 * 5, 120)
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)

   def forward(self, x):

    x = self.pool(F.relu(self.conv1(x)))
    x = self.pool(F.relu(self.conv2(x)))
    x = x.view(-1, 16 * 5 * 5)
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
    return x

Initialize model

model = TheModelClass()

Initialize optimizer

optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

Print model’s state_dict

print(“Model’s state_dict:”) for param_tensor in model.state_dict(): print(param_tensor, “\t”, model.state_dict()[param_tensor].size())

Print optimizer’s state_dict

print(“Optimizer’s state_dict:”) for var_name in optimizer.state_dict(): print(var_name, “\t”, optimizer.state_dict()[var_name])

上述代码先是简单定义一个 5 层的 CNN，然后分别打印模型的参数和优化器参数。<br />输出结果：
```python
Model's state_dict:
conv1.weight     torch.Size([6, 3, 5, 5])
conv1.bias   torch.Size([6])
conv2.weight     torch.Size([16, 6, 5, 5])
conv2.bias   torch.Size([16])
fc1.weight   torch.Size([120, 400])
fc1.bias     torch.Size([120])
fc2.weight   torch.Size([84, 120])
fc2.bias     torch.Size([84])
fc3.weight   torch.Size([10, 84])
fc3.bias     torch.Size([10])
Optimizer's state_dict:
state    {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]

2. 预测时加载和保存模型

加载/保存状态字典(推荐做法)

保存的代码：

torch.save(model.state_dict(), PATH)

加载的代码：

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

当需要为预测保存一个模型的时候，只需要保存训练模型的可学习参数即可。采用 torch.save() 来保存模型的状态字典的做法可以更方便加载模型，这也是推荐这种做法的原因。
通常会用 .pt 或者 .pth 后缀来保存模型。
记住

1. 在进行预测之前，必须调用 model.eval() 方法来将 dropout 和 batch normalization 层设置为验证模型。否则，只会生成前后不一致的预测结果。
2. load_state_dict() 方法必须传入一个字典对象，而不是对象的保存路径，也就是说必须先反序列化字典对象，然后再调用该方法，也是例子中先采用 torch.load() ，而不是直接 model.load_state_dict(PATH)

   加载/保存整个模型

   保存：

   torch.save(model, PATH)

   加载：

   # Model class must be defined somewhere
   model = torch.load(PATH)
   model.eval()

   保存和加载模型都是采用非常直观的语法并且都只需要几行代码即可实现。这种实现保存模型的做法将是采用 Python 的 pickle 模块来保存整个模型，这种做法的缺点就是序列化后的数据是属于特定的类和指定的字典结构，原因就是 pickle 并没有保存模型类别，而是保存一个包含该类的文件路径，因此，当在其他项目或者在 refactors 后采用都可能出现错误。

   3. 加载和保存的检查点（断点训练）

   保存的示例代码：

   torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

   加载的示例代码：
   ```python model = TheModelClass(args, **kwargs) optimizer = TheOptimizerClass(args, **kwargs)

checkpoint = torch.load(PATH) model.load_state_dict(checkpoint[‘model_state_dict’]) optimizer.load_state_dict(checkpoint[‘optimizer_state_dict’]) epoch = checkpoint[‘epoch’] loss = checkpoint[‘loss’]

model.eval()

- or -

model.train()

**当保存一个通用的检查点(checkpoint)时，无论是用于继续训练还是预测，都需要保存更多的信息**，不仅仅是 state_dict ，比如说优化器的 state_dict 也是非常重要的，它包含了用于模型训练时需要更新的参数和缓存信息，还可以保存的信息包括 epoch，即中断训练的批次，最后一次的训练 loss，额外的 torch.nn.Embedding 层等等。<br />上述保存代码就是介绍了如何保存这么多种信息，通过用一个字典来进行组织，然后继续调用 torch.save 方法，一般保存的文件后缀名是 .tar 。<br />加载代码也如上述代码所示，首先需要初始化模型和优化器，然后加载模型时分别调用 torch.load 加载对应的 state_dict 。然后通过不同的键来获取对应的数值。<br />**加载完后，根据后续步骤，调用 model.eval() 用于预测，model.train() 用于恢复训练。**
<a name="RosVN"></a>
### **4. 不同设备下保存和加载模型**
<a name="YRnS4"></a>
#### 在GPU上保存模型，在 CPU 上加载模型
```python
# 保存模型的示例代码：
torch.save(model.state_dict(), PATH)
# 加载模型的示例代码：
device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))

在 CPU 上加载在 GPU 上训练的模型，必须在调用 torch.load() 的时候，设置参数 map_location ，指定采用的设备是 torch.device(‘cpu’)，这个做法会将张量都重新映射到 CPU 上。

在GPU上保存模型，在 GPU 上加载模型

# 保存模型的示例代码：
torch.save(model.state_dict(), PATH)
# 加载模型的示例代码：
device = torch.device('cuda')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH)
model.to(device)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

在 GPU 上训练和加载模型，调用 torch.load() 加载模型后，还需要采用 model.to(torch.device(‘cuda’))，将模型调用到 GPU 上，并且后续输入的张量都需要确保是在 GPU 上使用的，即也需要采用 my_tensor.to(device)。

在CPU上保存，在GPU上加载模型

# 保存模型的示例代码：
torch.save(model.state_dict(), PATH)
# 加载模型的示例代码：
device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
model.to(device)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

保存 torch.nn.DataParallel 模型

# 保存模型的示例代码：
torch.save(model.module.state_dict(), PATH)

torch.nn.DataParallel 是用于实现多 GPU 并行的操作，保存模型的时候，是采用 model.module.state_dict()。
加载模型的代码也是一样的，采用 torch.load() ，并可以放到指定的 GPU 显卡上。

二、模型部署

1. TorchScript

pytorch本身是一个eager模式（ 命令式编程环境，可以立即评估操作产生的结果，无需构建计算图。 ）设计的深度学习框架，易于debug和观察，但是不利于性能的解耦和优化。PyTorch JIT模式就是一种将原本eager模式的表达转变并固定为一个计算图，便于进行优化和序列化。（向TF靠拢）

Tracing 方式

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)
    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h
my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell)
'''
output:
    MyCell(
      original_name=MyCell
      (linear): Linear(original_name=Linear)
    )
'''

TorchScript将其定义记录在一个中间表征（或IR）中，在深度学习中通常被称为图。我们可以用.graph属性检查该图。

print(traced_cell.graph)
'''
graph(%self.1 : __torch__.MyCell,
      %x : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu),
      %h : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu)):
  %linear : __torch__.torch.nn.modules.linear.Linear = prim::GetAttr[name="linear"](%self.1)
  %20 : Tensor = prim::CallMethod[name="forward"](%linear, %x)
  %11 : int = prim::Constant[value=1]() # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  %12 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::add(%20, %h, %11) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  %13 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::tanh(%12) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  %14 : (Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu), Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu)) = prim::TupleConstruct(%13, %13)
  return (%14)
'''

然而，这是一个非常低级的表示法，图中包含的大部分信息对终端用户没有用处。相反，我们可以使用.code属性来给出代码的Python语法解释。

print(traced_cell.code)
'''
def forward(self,
    x: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  linear = self.linear
  _0 = torch.tanh(torch.add((linear).forward(x, ), h))
  return (_0, _0)
'''

• TorchScript代码可以在自己的解释器中调用，这基本上是一个受限的Python解释器。这个解释器不会获得全局解释器锁，因此可以在同一个实例上同时处理许多请求。
• 这种格式允许我们将整个模型保存在磁盘上，并将其加载到另一个环境中，例如在用Python以外的语言编写的服务器中。
• TorchScript为我们提供了一种表示方法，我们可以对代码进行编译器优化，以提供更有效的执行。TorchScript允许我们与许多后端/设备运行机制对接，这些运行机制需要比单个运算符更广泛的程序视图。

我们可以看到，调用traced_cell产生的结果与Python模块相同。

print(my_cell(x, h))
print(traced_cell(x, h))
'''
(tensor([[0.9328, 0.7275, 0.3903, 0.8940],
        [0.9388, 0.4139, 0.7735, 0.8792],
        [0.8383, 0.7193, 0.6531, 0.9163]], grad_fn=<TanhBackward0>), tensor([[0.9328, 0.7275, 0.3903, 0.8940],
        [0.9388, 0.4139, 0.7735, 0.8792],
        [0.8383, 0.7193, 0.6531, 0.9163]], grad_fn=<TanhBackward0>))
(tensor([[0.9328, 0.7275, 0.3903, 0.8940],
        [0.9388, 0.4139, 0.7735, 0.8792],
        [0.8383, 0.7193, 0.6531, 0.9163]],
       grad_fn=<DifferentiableGraphBackward>), tensor([[0.9328, 0.7275, 0.3903, 0.8940],
        [0.9388, 0.4139, 0.7735, 0.8792],
        [0.8383, 0.7193, 0.6531, 0.9163]],
       grad_fn=<DifferentiableGraphBackward>))
'''

Scripting 方式

class MyDecisionGate(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x
class MyCell(torch.nn.Module):
    def __init__(self, dg):
        super(MyCell, self).__init__()
        self.dg = dg
        self.linear = torch.nn.Linear(4, 4)
    def forward(self, x, h):
        new_h = torch.tanh(self.dg(self.linear(x)) + h)
        return new_h, new_h
my_cell = MyCell(MyDecisionGate())
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell.dg.code)
print(traced_cell.code)
'''
def forward(self,
    argument_1: Tensor) -> Tensor:
  return torch.neg(argument_1)
def forward(self,
    x: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  dg = self.dg
  linear = self.linear
  _0 = torch.add((dg).forward((linear).forward(x, ), ), h)
  _1 = torch.tanh(_0)
  return (_1, _1)
  '''

看一下 traced_cell.code 输出，我们可以看到if-else分支无处可寻！为什么？追踪所做的正是我们所说的：运行代码，记录所发生的操作，并构建一个ScriptModule，正是这样做的。不幸的是，控制流之类的东西被抹去了。
我们怎样才能在TorchScript中忠实地表现这个模块呢？我们提供了一个脚本编译器，它对你的Python源代码进行直接分析，将其转化为TorchScript。让我们使用脚本编译器转换MyDecisionGate。

scripted_gate = torch.jit.script(MyDecisionGate())
my_cell = MyCell(scripted_gate)
scripted_cell = torch.jit.script(my_cell)
print(scripted_gate.code)
print(scripted_cell.code)
'''
def forward(self,
    x: Tensor) -> Tensor:
  if bool(torch.gt(torch.sum(x), 0)):
    _0 = x
  else:
    _0 = torch.neg(x)
  return _0
def forward(self,
    x: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  dg = self.dg
  linear = self.linear
  _0 = torch.add((dg).forward((linear).forward(x, ), ), h)
  new_h = torch.tanh(_0)
  return (new_h, new_h)
'''

Scripting和Tracing组合

有些情况下需要使用追踪而不是脚本（例如，一个模块有许多基于Python常量值的架构决定，我们希望这些决定不要出现在TorchScript中）。在这种情况下，脚本可以与追踪结合起来：torch.jit.script将内联追踪模块的代码，而追踪将内联脚本模块的代码。

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)
    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h
class MyRNNLoop(torch.nn.Module):
    def __init__(self):
        super(MyRNNLoop, self).__init__()
        self.cell = torch.jit.trace(MyCell(scripted_gate), (x, h))
    def forward(self, xs):
        h, y = torch.zeros(3, 4), torch.zeros(3, 4)
        for i in range(xs.size(0)):
            y, h = self.cell(xs[i], h)
        return y, h
rnn_loop = torch.jit.script(MyRNNLoop())
print(rnn_loop.code)
'''
def forward(self,
    xs: Tensor) -> Tuple[Tensor, Tensor]:
  h = torch.zeros([3, 4])
  y = torch.zeros([3, 4])
  y0 = y
  h0 = h
  for i in range(torch.size(xs, 0)):
    cell = self.cell
    _0 = (cell).forward(torch.select(xs, 0, i), h0, )
    y1, h1, = _0
    y0, h0 = y1, h1
  return (y0, h0)
'''

class WrapRNN(torch.nn.Module):
    def __init__(self):
        super(WrapRNN, self).__init__()
        self.loop = torch.jit.script(MyRNNLoop())
    def forward(self, xs):
        y, h = self.loop(xs)
        return torch.relu(y)
traced = torch.jit.trace(WrapRNN(), (torch.rand(10, 3, 4)))
print(traced.code)
'''
def forward(self,
    xs: Tensor) -> Tensor:
  loop = self.loop
  _0, h, = (loop).forward(xs, )
  return torch.relu(h)
  '''

保存与加载模型

我们提供 API 以存档格式在磁盘中保存和加载 TorchScript 模块。这种格式包括代码、参数、属性和调试信息，这意味着存档是模型的独立表示，可以在完全独立的过程中加载。让我们保存并加载我们包装好的 RNN 模块：

traced.save('wrapped_rnn.pt')
loaded = torch.jit.load('wrapped_rnn.pt')
print(loaded)
print(loaded.code)

trace和script的区别

1、trace只记录走过的tensor和对tensor的操作，不会记录任何控制流信息，如if条件句和循环。因为没有记录控制流的另外的路，也没办法对其进行优化。好处是trace深度嵌入python语言，复用了所有python的语法，在计算流中记录数据流。
2、script会去理解所有的code，真正像一个编译器一样去进行lexer、parser、Semantic analusis的分析「也就是词法分析语法分析句法分析，形成AST树，最后再将AST树线性化」。script相当于一个嵌入在Python/Pytorch的DSL，其语法只是pytorch语法的子集，这意味着存在一些op和语法script不支持，这样在编译的时候就会遇到问题。此外，script的编译优化方式更像是CPU上的传统编译优化，重点对于图进行硬件无关优化，并对IF、loop这样的statement进行优化。

例子

简单的 MNIST demo，从使用 Python 训练到用 JIT 将 Python 模型转换为 TorchScript Module，然后用 C++ 加载 TorchScript Module 做推断的完整的过程：

2. onnx

函数

torch.onnx.export(model, args, f, export_params=True, verbose=False,
                          training=False, input_names=None, output_names=None)

将一个模型导出到ONNX格式。该exporter会运行一次你的模型，以便于记录模型的执行轨迹，并将其导出；目前，exporter还不支持动态模型(例如，RNNs)。
参数

• model(torch.nn.Module)-要被导出的模型
• args(参数的集合)-模型的输入，例如，这种model(*args)方式是对模型的有效调用。任何非Variable参数都将硬编码到导出的模型中；任何Variable参数都将成为导出的模型的输入，并按照他们在args中出现的顺序输入。如果args是一个Variable，这等价于用包含这个Variable的1-ary元组调用它。(注意:现在不支持向模型传递关键字参数。)
• f-一个类文件的对象(必须实现文件描述符的返回)或一个包含文件名字符串。一个二进制Protobuf将会写入这个文件中。
• export_params(bool,default True)-如果指定，所有参数都会被导出。如果你只想导出一个未训练的模型，就将此参数设置为False。在这种情况下，导出的模型将首先把所有parameters作为参arguments，顺序由model.state_dict().values()指定。
• verbose(bool,default False)-如果指定，将会输出被导出的轨迹的调试描述。
• training(bool,default False)-导出训练模型下的模型。目前，ONNX只面向推断模型的导出，所以一般不需要将该项设置为True。
• input_names(list of strings, default empty list)-按顺序分配名称到图中的输入节点。
• output_names(list of strings, default empty list)-按顺序分配名称到图中的输出节点。

  Example: AlexNet from PyTorch to ONNX

  这是一个简单的脚本，它将预训练的 AlexNet 导出到名为 alexnet.onnx 的 ONNX 文件。对 torch.onnx.export 的调用运行模型一次以跟踪其执行，然后将跟踪的模型导出到指定文件： ```python import torch import torchvision

dummy_input = torch.randn(10, 3, 224, 224, device=”cuda”) model = torchvision.models.alexnet(pretrained=True).cuda()

Providing input and output names sets the display names for values

within the model’s graph. Setting these does not change the semantics

of the graph; it is only for readability.

#

The inputs to the network consist of the flat list of inputs (i.e.

the values you would pass to the forward() method) followed by the

flat list of parameters. You can partially specify names, i.e. provide

a list here shorter than the number of inputs to the model, and we will

only set that subset of names, starting from the beginning.

inputnames = [ “actual_input_1” ] + [ “learned%d” % i for i in range(16) ] output_names = [ “output1” ]

model.eval()# 若模型将用来推理或测试，则需要将模型设置为eval

torch.onnx.export(model, dummy_input, “alexnet.onnx”, verbose=True, input_names=input_names, output_names=output_names)

由此产生的alexnet.onnx文件包含一个二进制协议缓冲区，其中包含网络结构和你导出的模型（在本例中是AlexNet）的参数。参数verbose=True会使导出器打印出模型的人类可读表示。
```python
# These are the inputs and parameters to the network, which have taken on
# the names we specified earlier.
graph(%actual_input_1 : Float(10, 3, 224, 224)
      %learned_0 : Float(64, 3, 11, 11)
      %learned_1 : Float(64)
      %learned_2 : Float(192, 64, 5, 5)
      %learned_3 : Float(192)
      # ---- omitted for brevity ----
      %learned_14 : Float(1000, 4096)
      %learned_15 : Float(1000)) {
  # Every statement consists of some output tensors (and their types),
  # the operator to be run (with its attributes, e.g., kernels, strides,
  # etc.), its input tensors (%actual_input_1, %learned_0, %learned_1)
  %17 : Float(10, 64, 55, 55) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[11, 11], pads=[2, 2, 2, 2], strides=[4, 4]](%actual_input_1, %learned_0, %learned_1), scope: AlexNet/Sequential[features]/Conv2d[0]
  %18 : Float(10, 64, 55, 55) = onnx::Relu(%17), scope: AlexNet/Sequential[features]/ReLU[1]
  %19 : Float(10, 64, 27, 27) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%18), scope: AlexNet/Sequential[features]/MaxPool2d[2]
  # ---- omitted for brevity ----
  %29 : Float(10, 256, 6, 6) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%28), scope: AlexNet/Sequential[features]/MaxPool2d[12]
  # Dynamic means that the shape is not known. This may be because of a
  # limitation of our implementation (which we would like to fix in a
  # future release) or shapes which are truly dynamic.
  %30 : Dynamic = onnx::Shape(%29), scope: AlexNet
  %31 : Dynamic = onnx::Slice[axes=[0], ends=[1], starts=[0]](%30), scope: AlexNet
  %32 : Long() = onnx::Squeeze[axes=[0]](%31), scope: AlexNet
  %33 : Long() = onnx::Constant[value={9216}](), scope: AlexNet
  # ---- omitted for brevity ----
  %output1 : Float(10, 1000) = onnx::Gemm[alpha=1, beta=1, broadcast=1, transB=1](%45, %learned_14, %learned_15), scope: AlexNet/Sequential[classifier]/Linear[6]
  return (%output1);
}

您还可以使用 ONNX 库验证输出，您可以使用 conda 安装该库：

conda install -c conda-forge onnx

import onnx
# Load the ONNX model
model = onnx.load("alexnet.onnx")
# Check that the model is well formed
onnx.checker.check_model(model)
# Print a human readable representation of the graph
print(onnx.helper.printable_graph(model.graph))

您还可以使用支持 ONNX 的众多运行时之一运行导出的模型。例如，安装 ONNX Runtime 后，您可以加载并运行模型：

import onnxruntime as ort
ort_session = ort.InferenceSession("alexnet.onnx")
outputs = ort_session.run(
    None,
    {"actual_input_1": np.random.randn(10, 3, 224, 224).astype(np.float32)},
)
print(outputs[0])

使用np.testing.assert_allclose检查pytorch模型和导出的onnx模型的输出是否一致，如果不一致，会报error。

x = torch.randn(size=(batchsize, 3, h, w),dtype=torch.float32).to(self.device)
model.eval() # pytorch模型
torch.onnx.export(model, x, save_onnx_path,opset_version=11,input_names=['input'], output_names=['output'])
with torch.no_grad():
    torch_out = model(x)
import onnxruntime
ort_session = onnxruntime.InferenceSession(save_onnx_path)
# compute ONNX Runtime output prediction
ort_inputs = {ort_session.get_inputs()[0].name:x.cpu().numpy()}
ort_outs = ort_session.run(None, ort_inputs)
# compare ONNX Runtime and PyTorch results
np.testing.assert_allclose(torch_out.cpu().numpy(), ort_outs[0], rtol=1e-03, atol=1e-05)

在内部，torch.onnx.export() 需要 torch.jit.ScriptModule 而不是 torch.nn.Module。如果传入的模型还不是 ScriptModule，export() 将使用跟踪将其转换为：

• Tracing：如果用一个还不是 ScriptModule 的 Module 调用 torch.onnx.export()，它首先执行与 torch.jit.trace() 等效的操作，它使用给定的 args 执行一次模型并记录所有操作在执行期间发生。这意味着如果您的模型是动态的，例如根据输入数据改变行为，导出的模型将不会捕获这种动态行为。类似地，跟踪可能仅对特定输入大小有效。我们建议检查导出的模型并确保操作符看起来合理。跟踪将展开循环和 if 语句，导出与跟踪运行完全相同的静态图。如果要使用动态控制流导出模型，则需要使用Scripting。
• Scripting：通过脚本编译模型可以保留动态控制流，并且对不同大小的输入有效。使用脚本：

        - 使用 torch.jit.script() 生成一个 ScriptModule。
        - 以 ScriptModule 为模型调用 torch.onnx.export()，并设置 example_outputs 参数。这是必需的，以便可以在不执行模型的情况下捕获输出的类型和形状。

  陷阱

  PyTorch模型可以使用NumPy或Python类型和函数来编写，但在追踪过程中，任何NumPy或Python类型的变量（而不是torch.Tensor）都会被转换为常量，如果这些值应该根据输入而改变，就会产生错误的结果。 ```python

  错误！！！ Will be replaced with constants during tracing.

  x, y = np.random.rand(1, 2), np.random.rand(1, 2) np.concatenate((x, y), axis=1)

Good！！! Tensor operations will be captured during tracing.

x, y = torch.randn(1, 2), torch.randn(1, 2) torch.cat((x, y), dim=1)

另外，**onnxruntime推理时输入的是float32的numpy数据**。
<a name="QfMuw"></a>
### 3. caffe2
为了运行导出的caffe2版本的脚本，你需要以下两项支持：<br />    你需要安装caffe2。如果你还没有安装，请按照以下说明进行安装：[https://caffe2.ai/docs/getting-started.html](https://caffe2.ai/docs/getting-started.html)。<br />    你需要安装onnx-caffe2，一个纯Python库，它为ONNX提供了一个caffe2的编译器。你可以用pip安装onnx-caffe2：
```python
pip install onnx-caffe2

# ...continuing from above
import onnx_caffe2.backend as backend
import numpy as np
# or "CPU"
rep = backend.prepare(model, device="CUDA:0") 
# For the Caffe2 backend:
#     rep.predict_net is the Caffe2 protobuf for the network
#     rep.workspace is the Caffe2 workspace for the network
#       (see the class onnx_caffe2.backend.Workspace)
outputs = rep.run(np.random.randn(10, 3, 224, 224).astype(np.float32))
# To run networks with more than one input, pass a tuple
# rather than a single numpy ndarray.
print(outputs[0])

参考或转载：
有jit和script的源码解析
onnx和jit速度比较