1. PyTorch及ONNX环境准备

为了正常运行ONNX，我们需要安装最新的Pytorch

git clone --recursive https://github.com/pytorch/pytorchcd pytorchmkdir build && cd buildsudo cmake .. -DPYTHON_INCLUDE_DIR=/usr/include/python3.5  -DUSE_MPI=OFFmake installexport PYTHONPATH=$PYTHONPATH:/opt/pytorch/build

上面的"/opt/pytorch/build"是你前面build pytorch的目录，写对路径即可。

通过整个PyTorch的源码安装，PyTorch支持的相关ONNX库也会随之安装好。安装路径在：/usr/local/lib/python3.5/dist-packages/torch

运行如下命令安装ONNX的库：

conda install -c conda-forge onnx

此外，还需要安装onnx-caffe2，一个纯Python库，它为ONNX提供了一个caffe2的编译器。你可以用pip安装onnx-caffe2：

pip3 install onnx-caffe2

2. 准备好把PyTorch转换成ONNX的代码

在上面的pytorch2caffe2.py就是一段参考代码，把DeblurGAN训练好的模型转换成ONNX 。代码解释如下：

import osimport sysimport torchimport torch.onnximport torch.utils.model_zoofrom torch.autograd import Variablesys.path.append("../DeblurGAN")from models.models import create_modelimport models.networks as networksfrom options.test_options import TestOptionsimport shutilimport onnxfrom onnx_caffe2.backend import Caffe2Backendbatch_size = 1    # just a random number# Load the pretrained model weightsmodel_path = './model/char_deblur/latest_net_G.pth'onnx_model_path = "./deblurring.onnx.pb"state_dict = torch.utils.model_zoo.load_url(model_path, model_dir="./model/char_deblur")# Load the DeblurnGAN neural networkgan_opt = TestOptions().parse()gan_opt.name = "char_deblur"gan_opt.checkpoints_dir = "./model/"gan_opt.model = "test"gan_opt.dataset_mode = "single"gan_opt.dataroot = "/tmp/gan/"try:    shutil.rmtree(gan_opt.dataroot)except:    passos.mkdir(gan_opt.dataroot)gan_opt.loadSizeX = 64gan_opt.loadSizeY = 64gan_opt.fineSize = 64gan_opt.learn_residual = Truegan_opt.nThreads = 1  # test code only supports nThreads = 1gan_opt.batchSize = 1  # test code only supports batchSize = 1gan_opt.serial_batches = True  # no shufflegan_opt.no_flip = True  # no flip#torch_model = create_model(gan_opt)gpus = []torch_model = networks.define_G(gan_opt.input_nc, gan_opt.output_nc, gan_opt.ngf,                                gan_opt.which_model_netG, gan_opt.norm, not gan_opt.no_dropout, gpus, False,                                gan_opt.learn_residual)torch_model.load_state_dict(state_dict)#torch_model.load_state_dict(state_dict)# set the train mode to false since we will only run the forward pass.torch_model.train(False)# Input to the modelx = Variable(torch.randn(batch_size, 3, 60, 60), requires_grad=True)x = x.float()# Export the modeltorch_out = torch.onnx._export(torch_model,             # model being run                               x,                       # model input (or a tuple for multiple inputs)                               onnx_model_path, # where to save the model (can be a file or file-like object)                               verbose=True, export_params=True, training=False)      # store the trained parameter weights inside the model fileonnx_model = onnx.load(onnx_model_path)onnx.checker.check_model(onnx_model)model_name = onnx_model_path.replace('.onnx.pb','')init_net, predict_net = Caffe2Backend.onnx_graph_to_caffe2_net(onnx_model.graph, device="CUDA")with open(model_name + "_init.pb", "wb") as f:    f.write(init_net.SerializeToString())with open(model_name + "_predict.pb", "wb") as f:    f.write(predict_net.SerializeToString())

基于这个例子中，用户需要自己修改的部分有如下几个：