精华内容
下载资源
问答
  • 模型转换(pthonnx转caffemodel) 一、背景介绍 因为模型要用到前端,所以要将pytorch训练的模型转成caffemodel,之前LZ已经写了一片tensorflow转caffemodel的教程,今天就总结一篇pytorch转onnx转caffemodel的...

    模型转换(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
    

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

    展开全文
  • opencv的dnn模块部署yolov5模型,pt转pth转onnx 看了大佬写的用opencv的dnn模块做yolov5目标检测对于里面的方法进行了探讨 关于pt转pth 原作者在文章的后面跟新了pt转pth的代码,这里我再放一下链接链接:pt2pth ...

    opencv的dnn模块部署yolov5模型,pt转pth转onnx

    看了大佬写的用opencv的dnn模块做yolov5目标检测对于里面的方法进行了探讨

    关于pt转pth

    自己尝试了很久也遇到各种问题,这其中的坑就不说了
    torch.load()和torch.save()是大坑
    原作者在文章的后面跟新了pt转pth的代码,这里我再放一下链接链接:pt2pth
    提取码:4y6s

    import torch
    from collections import OrderedDict
    import pickle
    import os
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    if __name__ == '__main__':
        choices = ['yolov5s', 'yolov5l', 'yolov5m', 'yolov5x']//修改yolov5s为你自己模型的名字
        modelfile = choices[0] + '.pt'
        utl_model = torch.load(modelfile, map_location=device)
        utl_param = utl_model['model'].model
        torch.save(utl_param.state_dict(), os.path.splitext(modelfile)[0] + '.pth')
        own_state = utl_param.state_dict()
        print(len(own_state))
    
        numpy_param = OrderedDict()
        for name in own_state:
            numpy_param[name] = own_state[name].data.cpu().numpy()
        print(len(numpy_param))
        with open(os.path.splitext(modelfile)[0] + '_numpy_param.pkl', 'wb') as fw:
            pickle.dump(numpy_param, fw)
    

    关于pth转onnx

    大家可以去看原作者写的他的github上有,我用的1.0
    这里说一下我遇到的问题

    1. 首先是官方yolov5的版本
      我用的yolov5-V3.0版本
    2. 再pth导出到onnx时,需要配置相应的网络模型,需要把网络模型yaml文件转换成py文件这是作者写的yaml2py
      但是我遇到了很多问题,首先是导出的py文件无法被读取,这里需要自己再添加一个My_YOLO函数,修改最后的return
    from common import *
    
    class My_YOLO_backbone_head(nn.Module):
        def __init__(self, num_classes=12, anchors=[[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], training=False):
            super().__init__()
            self.seq0_Focus = Focus(3, 32, 3)
            self.seq1_Conv = Conv(32, 64, 3, 2)
            self.seq2_BottleneckCSP = BottleneckCSP(64, 64, 1)
            self.seq3_Conv = Conv(64, 128, 3, 2)
            self.seq4_BottleneckCSP = BottleneckCSP(128, 128, 3)
            self.seq5_Conv = Conv(128, 256, 3, 2)
            self.seq6_BottleneckCSP = BottleneckCSP(256, 256, 3)
            self.seq7_Conv = Conv(256, 512, 3, 2)
            self.seq8_SPP = SPP(512, 512, [5, 9, 13])
            self.seq9_BottleneckCSP = BottleneckCSP(512, 512, 1, False)
            self.seq10_Conv = Conv(512, 256, 1, 1)
            self.seq13_BottleneckCSP = BottleneckCSP(512, 256, 1, False)
            self.seq14_Conv = Conv(256, 128, 1, 1)
            self.seq17_BottleneckCSP = BottleneckCSP(256, 128, 1, False)
            self.seq18_Conv = Conv(128, 128, 3, 2)
            self.seq20_BottleneckCSP = BottleneckCSP(256, 256, 1, False)
            self.seq21_Conv = Conv(256, 256, 3, 2)
            self.seq23_BottleneckCSP = BottleneckCSP(512, 512, 1, False)
        def forward(self, x):
            x = self.seq0_Focus(x)
            x = self.seq1_Conv(x)
            x = self.seq2_BottleneckCSP(x)
            x = self.seq3_Conv(x)
            xRt0 = self.seq4_BottleneckCSP(x)
            x = self.seq5_Conv(xRt0)
            xRt1 = self.seq6_BottleneckCSP(x)
            x = self.seq7_Conv(xRt1)
            x = self.seq8_SPP(x)
            x = self.seq9_BottleneckCSP(x)
            xRt2 = self.seq10_Conv(x)
            route = F.interpolate(xRt2, size=(int(xRt2.shape[2] * 2), int(xRt2.shape[3] * 2)), mode='nearest')
            x = torch.cat([route, xRt1], dim=1)
            x = self.seq13_BottleneckCSP(x)
            xRt3 = self.seq14_Conv(x)
            route = F.interpolate(xRt3, size=(int(xRt3.shape[2] * 2), int(xRt3.shape[3] * 2)), mode='nearest')
            x = torch.cat([route, xRt0], dim=1)
            out0 = self.seq17_BottleneckCSP(x)/out0
            route = self.seq18_Conv(out0)//out0
            x = torch.cat([route, xRt3], dim=1)
            out1 = self.seq20_BottleneckCSP(x)//out1
            route = self.seq21_Conv(out1)//out1
            x = torch.cat([route, xRt2], dim=1)
            out2 = self.seq23_BottleneckCSP(x)//out2
            return out0, out1, out2
    //修改最后的return,按照上面的格式
    
    class My_YOLO(nn.Module)://这里需要自己添加
        def __init__(self, num_classes, anchors=(), training=False):
            super().__init__()
            self.backbone_head = My_YOLO_backbone_head()
            self.yolo_layers = Yolo_Layers(nc=num_classes, anchors=anchors, ch=(128,256,512),training=training)
        def forward(self, x):
            out0, out1, out2 = self.backbone_head(x)
            output = self.yolo_layers([out0, out1, out2])
            return output
    

    其次运行convert_onnx休要修改其中的类型名,改成自己模型的名字
    在这里插入图片描述
    这里就需要把自己模型结构的py文件名称替换yolov5s,这里就是上面需要My_YOLO函数的原因
    在这里插入图片描述

    说一下最后的结果

    模型成功导出,opencv的dnn模块也可以成功读取,但是经过2次模型转换,精度下降很多,不清楚是我自己的原因还是其他,检测效果很差
    希望有思路的同志交流一下

    下一篇;关于原文中,修改Focus类,把Detect类里面的1x1卷积定义在紧邻着Detect类之前的外面,然后去掉Detect类,组成新的model然后使用官方的export.py直接导出onnx

    参考:https://blog.csdn.net/nihate/article/details/112731327#comments_14884604

    展开全文
  • 2. 将pytorch训练的pth模型转换成ONNX,并编码测试; 3. 含训练和测试数据,含训练ok的pth模型和ONNX模型,含完整python和C++实现; 4. 使用方法:首先运行“TrainTestConvertOnnx.py”执行“训练数据读入、模型...
  • pthonnx转mnn bug总结

    2020-09-23 17:33:42
    pytorch版yolov3转onnx 样例 import torch import torchvision import numpy as np from onnxruntime.datasets import get_example import onnxruntime from onnx import shape_inference import onnx import os ...

    pytorch版yolov3转onnx

    样例

    import torch
    import torchvision
    import numpy as np
    from onnxruntime.datasets import get_example
    import onnxruntime
    from onnx import shape_inference
    import onnx
    import os
    from models import *
    
    img_size = 416
    cfg = 'cfg/yolov3.cfg'
    weights = "yolov3_models/yolov3.pt"
    
    model = Darknet(cfg, (img_size, img_size)).to("cpu")
    chkpt = torch.load(weights, map_location="cpu")
    
    chkpt['model'] = {k: v for k, v in chkpt['model'].items() if model.state_dict()[k].numel() == v.numel()}
    model.load_state_dict(chkpt['model'], strict=False)
    print('loaded weights from', weights, '\n')
    
    model.eval()
    
    x = torch.randn(1,3,416,416).cpu()
    export_onnx_file = "onnx_models/yolov3.onnx" 			
    torch.onnx.export(model,x,export_onnx_file,verbose=True, input_names=["input"], output_names=["output"])
    onnx.save(onnx.shape_inference.infer_shapes(onnx.load(export_onnx_file)), export_onnx_file)
    print("finish")
    def to_numpy(tensor):
        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
    
    model.eval()
    
    x = torch.randn(1,3,416,416).cpu()
    with torch.no_grad():
        torch_out = model(x)
    example_model = get_example(os.getcwd()+'/'+export_onnx_file)
    sess = onnxruntime.InferenceSession(example_model)
    onnx_out = sess.run(None, {sess.get_inputs()[0].name: to_numpy(x)})
    
    # print(torch_out[1].shape,torch_out[1][0][:5])
    # print(onnx_out[1].shape,onnx_out[1][0][:5])
    print(torch_out.shape,torch_out[0][0][:10])
    print(onnx_out[0].shape,onnx_out[0][0][0][:10])
    

    onnx非固定输入(动态输入和动态输出)

    import io
    import numpy as np
    
    from torch import nn
    import torch.utils.model_zoo as model_zoo
    import torch.onnx
    # Super Resolution model definition in PyTorch
    import torch.nn as nn
    import torch.nn.init as init
    import torch.nn.functional as F
    import onnxruntime
    import torchvision
    
    class SuperResolutionNet(nn.Module):
        def __init__(self, upscale_factor, inplace=False):
            super(SuperResolutionNet, self).__init__()
    
            self.relu = nn.ReLU(inplace=inplace)
            self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
            self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
            self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
            self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
            self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
            # self.upsample = nn.Upsample(scale_factor=upscale_factor, mode='nearest')
            self._initialize_weights()
    
        def forward(self, x):
            x = self.relu(self.conv1(x))
            x = self.relu(self.conv2(x))
            x = self.relu(self.conv3(x))  # torch.Size([1, 32, 224, 224])
            x = self.conv4(x)  # torch.Size([1, 9, 224, 224]) 224*3=672
            x = self.pixel_shuffle(x)  # torch.Size([1, 1, 672, 672])
            # x = self.upsample(x)
            x = F.interpolate(x, scale_factor=0.5)
            return x
    
        def _initialize_weights(self):
            init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
            init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
            init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
            init.orthogonal_(self.conv4.weight)
    
    
    # Create the super-resolution model by using the above model definition.
    torch_model = SuperResolutionNet(upscale_factor=3)
    
    # set the model to inference mode
    torch_model.eval()
    
    # Input to the model
    batch_size = 1  # just a random number
    channel = 1
    h_size = 224
    w_size = 224
    x = torch.randn(batch_size, channel, h_size, w_size, requires_grad=True)
    torch_out = torch_model(x)
    
    dynamic_axes = {'input': {0: 'batch_size', 1: 'channel', 2: "height", 3: 'width'},  # variable lenght axes
                    'output': {0: 'batch_size', 1: 'channel', 2: "height", 3: 'width'}}
    # Export the model
    torch.onnx.export(torch_model,  # model being run
                      x,  # model input (or a tuple for multiple inputs)
                      "super_resolution.onnx",  # 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=11,  # 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
                      dynamic_axes=dynamic_axes)
    def to_numpy(tensor):
        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
    
    ort_session = onnxruntime.InferenceSession("super_resolution.onnx")
    ort_x = torch.randn(4, 1, 666, 666, requires_grad=True)
    # compute ONNX Runtime output prediction
    ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(ort_x)}
    ort_outs = ort_session.run(None, ort_inputs)
    
    torch_out = torch_model(ort_x)
    # compare ONNX Runtime and PyTorch results
    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!")
    
    展开全文
  • 1.RuntimeError: Exporting the operator max_unpool2d to ONNX opset version 9 is not supported. 参考:https://www.wangt.cc/2020/10/runtimeerror-exporting-the-operator-max_unpool2d-to-onnx/ import ...

    1. RuntimeError: Exporting the operator max_unpool2d to ONNX opset version 9 is not supported. 

    参考:https://www.wangt.cc/2020/10/runtimeerror-exporting-the-operator-max_unpool2d-to-onnx/

    
    import argparse
    import logging
    import os
    
    import numpy as np
    import torch
    import torch.nn.functional as F
    from PIL import Image
    from torchvision import transforms
    
    from unet.ENet import ENet
    #from unet.CGNet import CGNet
    from utils.data_vis import plot_img_and_mask2
    from utils.USdataset import USDataset
    from utils.dataset import BasicDataset
    from torch.autograd import Variable
    
    
    
    if __name__ == "__main__":
    
    
        net = ENet(classes=1)
    
        # net = nn.DataParallel(net)
        # cudnn.benchmark = True
    
        # logging.info("Loading model {}".format(args.model))
    
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        logging.info(f'Using device {device}')
    
        if torch.cuda.device_count() >1:
            net = torch.nn.DataParallel(net)
        net.to(device=device)
    
        best_model=torch.load("CP_epoch30.pth", map_location=device)
        # print(best_model)
    
        net.load_state_dict(best_model,False)
    
    
        # model = UNet(n_channels=3, n_classes=4)
        # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        # 加载模型参数
        # model.load_state_dict(torch.load('best_model.pth', map_location=device))
        
        # An example input you would normally provide to your model's forward() method
        # x = Variable(torch.rand((1, 3, 896, 512),dtype=torch.float32))
    
        # x=np.random.ranint(((1, 3, 896, 512))
        img=torch.rand(1, 3, 896, 512)
    
        # img=torch.tensor(x)
    
        # img.to(device=device, dtype=torch.float32)
    
        img=img.to(device)
        # Export the model
        torch_out = torch.onnx._export(net, img, "enet.onnx",  export_params=True, verbose=True,operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK ,
                                                # verbose=True,  # store the trained parameter weights inside the model file
                                                # training=False
                                                input_names=['input'],
    					                        output_names=['output']
                                                )

     

    后tensorrt加载onnx转化engine

    fuck~~~~~~

    2 . 根据链接修改enet中索引池化为 :

    self.main_unpool1 =  nn.Upsample(scale_factor=2, mode='nearest')

    之后onnx再转trt出错如下:onnx found unsupported datatype (11) when importing initialize

    删除脚本中operator_export_type操作

    
    import argparse
    import logging
    import os
    
    import numpy as np
    import torch
    import torch.nn.functional as F
    from PIL import Image
    from torchvision import transforms
    
    from model.ENet import ENet
    
    from torch.autograd import Variable
    
    
    
    if __name__ == "__main__":
    
    
        net = ENet(classes=1)
    
        # net = nn.DataParallel(net)
        # cudnn.benchmark = True
    
        # logging.info("Loading model {}".format(args.model))
    
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        logging.info(f'Using device {device}')
    
        if torch.cuda.device_count() >1:
            net = torch.nn.DataParallel(net)
        net.to(device=device)
    
        # best_model=torch.load("best_model.pth", map_location=device)
        # net.load_state_dict(best_model,False)
    
        best_model=torch.load("best_model.pth", map_location=device)
        net.load_state_dict(best_model)
    
    
        img=torch.rand(1, 3, 320, 320)
    
      
    
        img=img.to(device)
    
    
        torch_out = torch.onnx._export(net, img, "enet.onnx", opset_version=11, export_params=True, verbose=True,
                                                # verbose=True,  # store the trained parameter weights inside the model file
                                                # training=False
                                                input_names=['input'],
    					                        output_names=['output']
                                                )

    之后使用trt自带的工具把上述生成的onnx转为trt

     .\trtexec.exe --onnx=enet.onnx --saveEngine=123.trt --device=0

    出现生成文件:

     

    参考:

    https://www.freesion.com/article/48831097285/

    https://gitee.com/antkillerfarm/onnx-tensorrt/blob/master/onnx2trt_utils.cpp

    https://github.com/onnx/onnx-tensorrt/tree/master/onnx_tensorrt

     

    最后发现转onnx时候模型没有冻结:

    
    import argparse
    import logging
    import os
    
    import numpy as np
    import torch
    import torch.nn.functional as F
    from PIL import Image
    from torchvision import transforms
    
    from model.ENet import ENet
    
    from torch.autograd import Variable
    
    
    
    if __name__ == "__main__":
    
    
        net = ENet(classes=1)
    
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
      
    
        if torch.cuda.device_count() >1:
            net = torch.nn.DataParallel(net)
        net.to(device=device)
    
        best_model=torch.load("CP_epoch30.pth", map_location=device)
        net.load_state_dict(best_model,False)
    
    
        net.eval()
    
        img=torch.rand(1, 3,896, 512)
    
        # img=torch.tensor(x)
    
        # img.to(device=device, dtype=torch.float32)
    
        img=img.to(device, dtype=torch.float32)
        # Export the model
    
    
        torch_out = torch.onnx._export(net, img, "CP_epoch30.onnx", opset_version=11, export_params=True, verbose=True,
                                                # verbose=True,  # store the trained parameter weights inside the model file
                                                # training=False
                                                input_names=['input'],
    					                        output_names=['output']
                                                )

     
    展开全文
  • 文章目录 pb转uff具体代码 运行结果 pb转onnx转trt具体代码 pb转onnx运行结果 onnx转化trt 方法1:trtexec 方法2:onnx-tensorrt工具 推理trt模型 pth转onnx onnx转pb 方法1:onnx-tensorflow工具 方法2:代码执行 pb转...
  • 文章目录引言基础概念onnx:跨框架的模型表达标准onnxruntime:部署模型的推理引擎示例代码0)安装onnxonnxruntime1)pytorch模型转onnx模型2)onnx模型检验3)调用ONNX Runtime测试输入图片参考教程 引言 目前...

空空如也

空空如也

1 2 3 4 5 ... 20
收藏数 817
精华内容 326
关键字:

pth转onnx