PyTorch 1.0 中文官方教程:使用ONNX将模型从PyTorch传输到Caffe2和移动端

2022-05-07 14:04:02 浏览数 (1)

译者:冯宝宝

在本教程中,我们将介绍如何使用ONNX将PyTorch中定义的模型转换为ONNX格式,然后将其加载到Caffe2中。一旦进入Caffe2,我们就可以运行模型来仔细检查它是否正确导出,然后我们展示了如何使用Caffe2功能(如移动导出器)在移动设备上执行模型。

在本教程中,你需要安装onnx和Caffe2。您可以使用pip install onnx获取onnx的二进制版本。

注意: 本教程需要PyTorch master分支,可以按照 这里说明进行安装。

代码语言:javascript复制
# 一些包的导入
import io
import numpy as np

from torch import nn
import torch.utils.model_zoo as model_zoo
import torch.onnx

超分辨率是一种提高图像,视频分辨率的方法,广泛用于图像处理或视频剪辑。在本教程中,我们将首先使用带有虚拟输入的小型超分辨率模型。

首先,让我们在PyTorch中创建一个SuperResolution模型。这个模型 直接来自PyTorch的例子,没有修改:

代码语言:javascript复制
# PyTorch中定义的Super Resolution模型
import torch.nn as nn
import torch.nn.init as init

class SuperResolutionNet(nn.Module):
    def __init__(self, upscale_factor, inplace=False):
        super(SuperResolutionNet, self).__init__()

阅读全文/改进本文

0 人点赞