pytorch转paddle推理模型

最近有个需求，是将训练好的pytorch模型转成paddlepaddle的inference_model，然后直接使用paddlepaddle载入使用。转换的工具主要使用paddle官方提供的X2paddle，对应项目链接：
https://github.com/PaddlePaddle/X2Paddle

官方文档中有对应pytorch模型转paddlepaddle模型的教程，但我只需要inference_model，所以我采用的方法是先将训练好的pytorch模型转成ONNX格式，然后在用X2Paddle将ONNX模型转成inference_model。

---

Pytorch转ONNX

Pytorch转ONNX的官方文档：https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html

这里默认已经将模型训练好了，那么可以使用以下脚本进行转换。下面示例中是我训练好的mobile_v3_small模型，model是实例化好的模型，model_weight_path是要载入的训练好的权重路径，onnx_file_name是保存的ONNX模型的名称，还需要注意的是测试使用的变量x由于我传入的是尺寸为224x224的RGB图片所以是torch.rand(1, 3, 224, 224, requires_grad=True, device=device)这里也需要根据自己的实际情况修改。

from PIL import Image
import torchvision.transforms as transforms
import torch
import torch.onnx
import onnx
import onnxruntime
import numpy as np
from model_v3 import mobilenet_v3_small

device = torch.device("cpu")


def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()


def main(save_path=None):
    assert isinstance(save_path, str), "lack of save_path parameter..."
    # create model
    model = mobilenet_v3_small(num_classes=4).to(device)
    # load model weights
    model_weight_path = "./weights/model-9.pth"
    model.load_state_dict(torch.load(model_weight_path, map_location=device))
    model.eval()
    # input to the model
    # [batch, channel, height, width]
    x = torch.rand(1, 3, 224, 224, requires_grad=True, device=device)
    torch_out = model(x)

    # export the model
    torch.onnx.export(model,                       # model being run
                      x,                           # model input (or a tuple for multiple inputs)
                      save_path,                   # where to save the model (can be a file or file-like object)
                      export_params=True,          # store the trained parameter weights inside the model file
                      opset_version=9,            # the ONNX version to export the model to
                      do_constant_folding=True,    # whether to execute constant folding for optimization
                      input_names=["input"],       # the model's input names
                      output_names=["output"],     # the model's output names
                      )

    # check onnx model
    onnx_model = onnx.load(save_path)
    onnx.checker.check_model(onnx_model)

    ort_session = onnxruntime.InferenceSession(save_path)

    # compute ONNX Runtime output prediction
    ort_inputs = {
    
    ort_session.get_inputs()[0].name: to_numpy(x)}
    ort_outs = ort_session.run(None, ort_inputs)

    # compare ONNX Runtime and Pytorch results
    # assert_allclose: Raises an AssertionError if two objects are not equal up to desired tolerance.
    np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)
    print("Exported model has been tested with ONNXRuntime, and the result looks good!")

    # load test image
    # img = Image.open("1.jpg")
    # 
    # # pre-process
    # preprocess = transforms.Compose([transforms.Resize([224, 224]),
    #                                  transforms.ToTensor(),
    #                                  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    # img = preprocess(img)
    # img = img.unsqueeze_(0)
    # 
    # # feed image into onnx model
    # ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(img)}
    # ort_outs = ort_session.run(None, ort_inputs)
    # prediction = ort_outs[0]
    # 
    # # np softmax process
    # prediction -= np.max(prediction, keepdims=True)  # 为了稳定地计算softmax概率， 一般会减掉最大元素
    # prediction = np.exp(prediction) / np.sum(np.exp(prediction), keepdims=True)
    # print(prediction)


if __name__ == '__main__':
    onnx_file_name = "mobilenetv3.onnx"
    main(save_path=onnx_file_name)

执行转换脚本后，只要程序没有报错并打印如下信息就说明转换成功：

Exported model has been tested with ONNXRuntime, and the result looks good!

---

ONNX转Paddle inference_model

安装X2Paddle，官方提供了两种安装方法，一种是pip安装，一种是源码安装，这里以pip安装为例：

pip install x2paddle

安装完后在刚刚生成的ONNX模型的目录下打开终端，然后输入下面指令即可转换，其中--model参数是指向刚刚生成的ONNX文件，记得改成自己ONNX文件名称，save_dir参数是转换后保存的地址：

x2paddle --framework=onnx --model=onnx_model.onnx --save_dir=pd_model

下面是我转化过程中的终端信息，DeprecationWarning警告可以忽略：

$ x2paddle --framework=onnx --model=mobilenetv3.onnx --save_dir=pd_model
paddle.__version__ = 2.1.0
Now translating model from onnx to paddle.
model ir_version: 6, op version: 9
shape inferencing ...
shape inferenced.
Now, onnx2paddle support convert onnx model opset_verison [9],opset_verison of your onnx model is 9, automatically treated as op_set: 9.
Total nodes: 231
Nodes converting ...
Converting node 498 ...     
Nodes converted.
Exporting inference model from python code ('/home/wz/my_project/mobilenet_v3/pd_model/x2paddle_code.py')... 

/home/wz/miniconda3/envs/paddle/lib/python3.7/site-packages/paddle/fluid/layers/utils.py:77: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated since Python 3.3,and in 3.9 it will stop working
  return (isinstance(seq, collections.Sequence) and

转换后的输出目录如下：

├── pd_model 转换文件生成目录
│         ├── inference_model  
│         │     ├── model.pdiparams
│         │     ├── model.pdiparams.info
│         │     └── model.pdmodel
│         ├── x2paddle_code.py
│         └── model.pdparams

其中inference_model文件夹就是我需要的推理模型。x2paddle_code.py是转换后paddlepaddle构建的模型文件。model.pdparams是对应模型文件的权重（不是推理模型权重）。

---

使用Paddle Inference预测

下面就用刚刚转换生成的inference_model文件来进行预测（主要是model.pdmodel模型文件以及model.pdiparams权重文件）。下面的代码主要参考paddle官方的demo：
https://github.com/PaddlePaddle/Paddle-Inference-Demo/

主要步骤：

• 初始化模型
• 读取预测图片并进行预处理（注意预处理方式要和你之前的预处理方式保持一致）
• 将图片输入网络进行预测

import os
from PIL import Image
import numpy as np

from paddle.inference import Config
from paddle.inference import create_predictor


def preprocess(img: Image.Image):
    img = img.resize((224, 224), resample=2)
    img = np.array(img, dtype=np.float32) / 255.
    img -= [0.485, 0.456, 0.406]
    img /= [0.229, 0.224, 0.225]
    img = img.transpose((2, 0, 1))  # HWC -> CHW
    img = img[np.newaxis, :]
    return img


def init_predictor(model_dir):
    model_path = os.path.join(model_dir, "model.pdmodel")
    params_path = os.path.join(model_dir, "model.pdiparams")
    config = Config(model_path, params_path)
    config.enable_memory_optim()

    # If not specific mkldnn, you can set the blas thread.
    # The thread num should not be greater than the number of cores in the CPU.
    config.set_cpu_math_library_num_threads(4)
    config.enable_mkldnn()

    predictor = create_predictor(config)
    return predictor


def run(pred, imgs):
    # copy img data to input tensor
    input_names = pred.get_input_names()
    for i, name in enumerate(input_names):
        input_tensor = pred.get_input_handle(name)
        input_tensor.reshape(imgs[i].shape)
        input_tensor.copy_from_cpu(imgs[i])

    # do the inference
    pred.run()

    results = []
    # get out data from output tensor
    output_names = pred.get_output_names()
    for i, name in enumerate(output_names):
        output_tensor = pred.get_output_handle(name)
        output_data = output_tensor.copy_to_cpu()
        results.append(output_data)

    return results


def main():
    pred = init_predictor("./pd_model/inference_model")
    img = Image.open("./1.jpg")
    img = preprocess(img)
    results = run(pred, [img])
    print(results[0])


if __name__ == '__main__':
    main()

下面是自己做测试时使用同一张图片分别使用Pytorch以及PaddlePaddle推理的结果：

# pytorch inference
[[-2.210124  -5.163459 -5.193339 12.506434 ]]

# paddle inference
[[-2.210112  -5.1634603 -5.1933465 12.506434 ]]