模型转换（pth转onnx转caffemodel）

一、背景介绍

因为模型要用到前端，所以要将pytorch训练的模型转成caffemodel，之前LZ已经写了一片tensorflow转caffemodel的教程，今天就总结一篇pytorch转onnx转caffemodel的教程

二、pth转onnx

这个严格意义上是DBFace系列的第四部分，但是主要是用来转模型，对其他模型也是适用的，于是就单独拎出来了，首先放上代码

import common
import eval_tool
import torch
import torch.nn as nn
import torch.nn.functional as F
# from dbface import DBFace
from dbface_mv2 import DBFace

import torch.onnx.symbolic_opset11
# import torch.onnx.symbolic_helper

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "3"

# trial_name = "small-H-dense-wide64-UCBA-keep12-noext-ignoresmall2"

trial_name = "mv2-320x320-without-wf_20200811"
jobdir = f"jobs/{trial_name}"


class OnnxModel(nn.Module):
    def __init__(self, **kwargs):
        super(OnnxModel, self).__init__()
		# 定义的模型结构
        self.model = DBFace(**kwargs)
        # 训练得到的模型结构和参数
        self.model.load(f"{jobdir}/models/74.pth")

    def forward(self, x):
        center, box, landmark = self.model(x)
        center_sigmoid = torch.sigmoid(center) # 当然也可以增加一些operation操作，这样后处理的时候也可以放到nnie上去做
        # center_maxpool = F.max_pool2d(center_sigmoid, kernel_size=3, padding=1, stride=1)
        # box = torch.exp(box)
        # return center_maxpool, box, landmark
        return center_sigmoid, box, landmark

#先要下载保存对应的模型
model = OnnxModel(has_landmark=True, wide=64, has_ext=True, upmode="DeCBA")
# 模型设置模式
model.eval()
model.cuda()

onnx_name = "mv2-288x480-without-exp"
common.mkdirs_from_file_path(f"{jobdir}/{onnx_name}.onnx")

# 模型输入大小
dummy = torch.zeros((1, 3, 288, 480)).cuda()
torch.onnx.export(model, dummy, f"{jobdir}/{onnx_name}.onnx", export_params=True, verbose=True,
                  output_names=["hm", "tlrb", "landmark"])

export的参数设置

def export(model, args, f, export_params=True, verbose=False, training=False,
           input_names=None, output_names=None, aten=False, export_raw_ir=False,
           operator_export_type=None, opset_version=None, _retain_param_name=True,
           do_constant_folding=False, example_outputs=None, strip_doc_string=True,
           dynamic_axes=None, keep_initializers_as_inputs=None):
    r"""
    Export a model into ONNX format.  This exporter runs your model
    once in order to get a trace of its execution to be exported;
    at the moment, it supports a limited set of dynamic models (e.g., RNNs.)
    See also: :ref:`onnx-export`
    Arguments:
        model (torch.nn.Module): the model to be exported.
        args (tuple of arguments): the inputs to
            the model, e.g., such that ``model(*args)`` is a valid
            invocation of the model.  Any non-Tensor arguments will
            be hard-coded into the exported model; any Tensor arguments
            will become inputs of the exported model, in the order they
            occur in args.  If args is a Tensor, this is equivalent
            to having called it with a 1-ary tuple of that Tensor.
            (Note: passing keyword arguments to the model is not currently
            supported.  Give us a shout if you need it.)
        f: a file-like object (has to implement fileno that returns a file descriptor)
            or a string containing a file name.  A binary Protobuf will be written
            to this file.
        export_params (bool, default True): if specified, all parameters will
            be exported.  Set this to False if you want to export an untrained model.
            In this case, the exported model will first take all of its parameters
            as arguments, the ordering as specified by ``model.state_dict().values()``
        verbose (bool, default False): if specified, we will print out a debug
            description of the trace being exported.
        training (bool, default False): export the model in training mode.  At
            the moment, ONNX is oriented towards exporting models for inference
            only, so you will generally not need to set this to True.
        input_names(list of strings, default empty list): names to assign to the
            input nodes of the graph, in order
        output_names(list of strings, default empty list): names to assign to the
            output nodes of the graph, in order
        aten (bool, default False): [DEPRECATED. use operator_export_type] export the
            model in aten mode. If using aten mode, all the ops original exported
            by the functions in symbolic_opset<version>.py are exported as ATen ops.
        export_raw_ir (bool, default False): [DEPRECATED. use operator_export_type]
            export the internal IR directly instead of converting it to ONNX ops.
        operator_export_type (enum, default OperatorExportTypes.ONNX):
            OperatorExportTypes.ONNX: all ops are exported as regular ONNX ops.
            OperatorExportTypes.ONNX_ATEN: all ops are exported as ATen ops.
            OperatorExportTypes.ONNX_ATEN_FALLBACK: if symbolic is missing,
                                                    fall back on ATen op.
            OperatorExportTypes.RAW: export raw ir.
        opset_version (int, default is 9): by default we export the model to the
            opset version of the onnx submodule. Since ONNX's latest opset may
            evolve before next stable release, by default we export to one stable
            opset version. Right now, supported stable opset version is 9.
            The opset_version must be _onnx_master_opset or in _onnx_stable_opsets
            which are defined in torch/onnx/symbolic_helper.py
        do_constant_folding (bool, default False): If True, the constant-folding
            optimization is applied to the model during export. Constant-folding
            optimization will replace some of the ops that have all constant
            inputs, with pre-computed constant nodes.
        example_outputs (tuple of Tensors, default None): example_outputs must be provided
            when exporting a ScriptModule or TorchScript Function.
        strip_doc_string (bool, default True): if True, strips the field
            "doc_string" from the exported model, which information about the stack
            trace.
        example_outputs: example outputs of the model that is being exported.
        dynamic_axes (dict<string, dict<int, string>> or dict<string, list(int)>, default empty dict):
            a dictionary to specify dynamic axes of input/output, such that:
            - KEY:  input and/or output names
            - VALUE: index of dynamic axes for given key and potentially the name to be used for
            exported dynamic axes. In general the value is defined according to one of the following
            ways or a combination of both:
            (1). A list of integers specifiying the dynamic axes of provided input. In this scenario
            automated names will be generated and applied to dynamic axes of provided input/output
            during export.
            OR (2). An inner dictionary that specifies a mapping FROM the index of dynamic axis in
            corresponding input/output TO the name that is desired to be applied on such axis of
            such input/output during export.
            Example. if we have the following shape for inputs and outputs:
                shape(input_1) = ('b', 3, 'w', 'h')
                and shape(input_2) = ('b', 4)
                and shape(output)  = ('b', 'd', 5)

            Then dynamic axes can be defined either as:
                (a). ONLY INDICES:
                    dynamic_axes = {'input_1':[0, 2, 3], 'input_2':[0], 'output':[0, 1]}

                    where automatic names will be generated for exported dynamic axes

                (b). INDICES WITH CORRESPONDING NAMES:
                    dynamic_axes = {'input_1':{0:'batch', 1:'width', 2:'height'},
                    'input_2':{0:'batch'},
                    'output':{0:'batch', 1:'detections'}

                    where provided names will be applied to exported dynamic axes

                (c). MIXED MODE OF (a) and (b)
                    dynamic_axes = {'input_1':[0, 2, 3], 'input_2':{0:'batch'}, 'output':[0,1]}
        keep_initializers_as_inputs (bool, default None): If True, all the initializers
            (typically corresponding to parameters) in the exported graph will also be
            added as inputs to the graph. If False, then initializers are not added as
            inputs to the graph, and only the non-parameter inputs are added as inputs.
            This may allow for better optimizations (such as constant folding etc.) by
            backends/runtimes that execute these graphs. If unspecified (default None),
            then the behavior is chosen automatically as follows. If operator_export_type
            is OperatorExportTypes.ONNX, the behavior is equivalent to setting this
            argument to False. For other values of operator_export_type, the behavior is
            equivalent to setting this argument to True.
    """

这一步其实还算比较简单的，当然还有一个需要考虑的问题，就是如果我要让onnx支持多batchsize怎么办呢？

这步LZ没有实测过，用在前端caffemodel，batchsize设定基本为1,所以多batchsize的使用场景可能用在后端居多。

inputs = ["input.1"]
outputs = ["hm", "tlrb", "landmark"]
dynamic_axes = {'input.1':{0:'batch'}, "hm":{0:'batch'}, "tlrb":{0:'batch'}, "landmark":{0:'batch'}}
torch.onnx.export(model, dummy, f"{jobdir}/{onnx_name}.onnx", export_params=True, verbose=True,
                  input_names =inputs, output_names=outputs，dynamic_axes=dynamic_axes)

三、onnx转caffemodel

主要是convertCaffe.py, github地址在https://github.com/MTlab/onnx2caffe

from __future__ import print_function
import sys
import caffe
import onnx
import numpy as np
from caffe.proto import caffe_pb2

caffe.set_mode_cpu()
from onnx2caffe._transformers import ConvAddFuser, ConstantsToInitializers
from onnx2caffe._graph import Graph

import onnx2caffe._operators as cvt
import onnx2caffe._weightloader as wlr
from onnx2caffe._error_utils import ErrorHandling
from collections import OrderedDict
from onnx import shape_inference
import importlib

transformers = [
    ConstantsToInitializers(),
    ConvAddFuser(),
]


def convertToCaffe(graph, prototxt_save_path, caffe_model_save_path):
    exist_edges = []
    layers = []
    exist_nodes = []
    err = ErrorHandling()
    for i in graph.inputs:
        edge_name = i[0]
        input_layer = cvt.make_input(i)
        layers.append(input_layer)
        exist_edges.append(i[0])
        graph.channel_dims[edge_name] = graph.shape_dict[edge_name][1]

    for id, node in enumerate(graph.nodes):
        node_name = node.name
        op_type = node.op_type
        inputs = node.inputs
        inputs_tensor = node.input_tensors
        input_non_exist_flag = False

        for inp in inputs:
            if inp not in exist_edges and inp not in inputs_tensor:
                input_non_exist_flag = True
                break
        if input_non_exist_flag:
            continue

        if op_type not in cvt._ONNX_NODE_REGISTRY:
            err.unsupported_op(node)
            continue
        converter_fn = cvt._ONNX_NODE_REGISTRY[op_type]
        layer = converter_fn(node, graph, err)
        if type(layer) == tuple:
            for l in layer:
                layers.append(l)
        else:
            layers.append(layer)
        outs = node.outputs
        for out in outs:
            exist_edges.append(out)

    net = caffe_pb2.NetParameter()
    for id, layer in enumerate(layers):
        layers[id] = layer._to_proto()
    net.layer.extend(layers)

    with open(prototxt_save_path, 'w') as f:
        print(net, file=f)

    caffe.set_mode_cpu()
    deploy = prototxt_save_path
    net = caffe.Net(deploy,
                    caffe.TEST)

    for id, node in enumerate(graph.nodes):
        node_name = node.name
        op_type = node.op_type
        inputs = node.inputs
        inputs_tensor = node.input_tensors
        input_non_exist_flag = False
        if op_type not in wlr._ONNX_NODE_REGISTRY:
            err.unsupported_op(node)
            continue
        converter_fn = wlr._ONNX_NODE_REGISTRY[op_type]
        converter_fn(net, node, graph, err)

    net.save(caffe_model_save_path)
    print("saving ", caffe_model_save_path)
    return net


def getGraph(onnx_path):
    model = onnx.load(onnx_path)
    model = shape_inference.infer_shapes(model)
    model_graph = model.graph
    graph = Graph.from_onnx(model_graph)
    graph = graph.transformed(transformers)
    graph.channel_dims = {}

    return graph


if __name__ == "__main__":
    onnx_path = sys.argv[1]
    prototxt_path = sys.argv[2]
    caffemodel_path = sys.argv[3]
    graph = getGraph(onnx_path)
    convertToCaffe(graph, prototxt_path, caffemodel_path)

这个转换LZ没有时间细看代码了，基本上支持的操作如下：

_ONNX_NODE_REGISTRY = {
    "Conv": _convert_conv,
    "Relu": _convert_relu,
    "BatchNormalization": _convert_BatchNorm,
    "Add": _convert_Add,
    "Mul": _convert_Mul,
    "Reshape": _convert_Reshape,
    "MaxPool": _convert_pool,
    "AveragePool": _convert_pool,
    "Dropout": _convert_dropout,
    "Gemm": _convert_gemm,
    "Upsample": _convert_upsample,
    "Concat": _convert_concat,
    "ConvTranspose": _convert_conv_transpose,
    "Sigmoid": _convert_sigmoid,
    "Flatten": _convert_Flatten,
}

写个脚本，运行对应的代码就行了

workspace="your folder path"
filename="name.onnx"
python convertCaffe.py $workspace$filename.onnx $workspace$filename.prototxt $workspace$filename.caffemodel

通过上述操作就可以完成这一套流程了。