• Python-ONNX 相关


    一、pytorch 转 onnx 推理加速

    01

    配置

    Ubuntu 16.04
    python 3.6
    onnx 1.6
    pytorch 1.5
    pycuda 2019.1.2
    torchvision 0.1.8

    建议详读,先安装好环境:

    https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#import_onnx_python

    02

    步骤

    1. 将pytorch模型转换成onnx模型

    这边用的是Darknet生成的pytoch模型

    import torchfrom torch.autograd import Variableimport onnx
    
    input_name = ['input']output_name = ['output']input = Variable(torch.randn(1, 3, 544, 544)).cuda()model = x.model.cuda()#x.model为我生成的模型
    # model = torch.load('', map_location="cuda:0")torch.onnx.export(model, input'model.onnx', input_names=input_name, output_names=output_name, verbose=True)

    其中

    #model = x.model.cuda()#若是不添加cuda()model = x.model

    出现报错

    RuntimeErrorInput type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same
     
    

    2. 检查模型

    model = onnx.load("model.onnx")onnx.checker.check_model(model)print("==> Passed")

    3. 测试onnx模型使用tensorrt推理前后对比

    import pycuda.autoinitimport numpy as npimport pycuda.driver as cudaimport tensorrt as trtimport torchimport osimport timefrom PIL import Imageimport cv2import torchvision
    filename = '000000.jpg'max_batch_size = 1onnx_model_path = 'yolo.onnx'
    TRT_LOGGER = trt.Logger()  # This logger is required to build an engine
    
    def get_img_np_nchw(filename):    image = cv2.imread(filename)    image_cv = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)    image_cv = cv2.resize(image_cv, (1920, 1080))    miu = np.array([0.485, 0.456, 0.406])    std = np.array([0.229, 0.224, 0.225])    img_np = np.array(image_cv, dtype=float) / 255.    r = (img_np[:, :, 0] - miu[0]) / std[0]    g = (img_np[:, :, 1] - miu[1]) / std[1]    b = (img_np[:, :, 2] - miu[2]) / std[2]    img_np_t = np.array([r, g, b])    img_np_nchw = np.expand_dims(img_np_t, axis=0)    return img_np_nchw
    class HostDeviceMem(object):    def __init__(self, host_mem, device_mem):        """Within this context, host_mom means the cpu memory and device means the GPU memory        """        self.host = host_mem        self.device = device_mem
        def __str__(self):        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
        def __repr__(self):        return self.__str__()
    
    def allocate_buffers(engine):    inputs = []    outputs = []    bindings = []    stream = cuda.Stream()    for binding in engine:        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size        dtype = trt.nptype(engine.get_binding_dtype(binding))        # Allocate host and device buffers        host_mem = cuda.pagelocked_empty(size, dtype)        device_mem = cuda.mem_alloc(host_mem.nbytes)        # Append the device buffer to device bindings.        bindings.append(int(device_mem))        # Append to the appropriate list.        if engine.binding_is_input(binding):            inputs.append(HostDeviceMem(host_mem, device_mem))        else:            outputs.append(HostDeviceMem(host_mem, device_mem))    return inputs, outputs, bindings, stream
    
    def get_engine(max_batch_size=1, onnx_file_path="", engine_file_path="", \               fp16_mode=False, int8_mode=False, save_engine=False,               ):    """Attempts to load a serialized engine if available, otherwise builds a new TensorRT engine and saves it."""
        def build_engine(max_batch_size, save_engine):        """Takes an ONNX file and creates a TensorRT engine to run inference with"""        EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)        with trt.Builder(TRT_LOGGER) as builder, \                builder.create_network(EXPLICIT_BATCH) as network, \                trt.OnnxParser(network, TRT_LOGGER) as parser:
                builder.max_workspace_size = 1 << 30  # Your workspace size            builder.max_batch_size = max_batch_size            # pdb.set_trace()            builder.fp16_mode = fp16_mode  # Default: False            builder.int8_mode = int8_mode  # Default: False            if int8_mode:                # To be updated                raise NotImplementedError
                # Parse model file            if not os.path.exists(onnx_file_path):                quit('ONNX file {} not found'.format(onnx_file_path))
                print('Loading ONNX file from path {}...'.format(onnx_file_path))            with open(onnx_file_path, 'rb') as model:                print('Beginning ONNX file parsing')                parser.parse(model.read())
                    if not parser.parse(model.read()):                    for error in range(parser.num_errors):                        print(parser.get_error(error))                    print("===========Parsing fail!!!!=================")                else :                    print('Completed parsing of ONNX file')
                print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))
                engine = builder.build_cuda_engine(network)            print("Completed creating Engine")
                if save_engine:                with open(engine_file_path, "wb") as f:                    f.write(engine.serialize())            return engine
        if os.path.exists(engine_file_path):        # If a serialized engine exists, load it instead of building a new one.        print("Reading engine from file {}".format(engine_file_path))        with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:            return runtime.deserialize_cuda_engine(f.read())    else:        return build_engine(max_batch_size, save_engine)
    
    def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):    # Transfer data from CPU to the GPU.    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]    # Run inference.    context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)    # Transfer predictions back from the GPU.    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]    # Synchronize the stream    stream.synchronize()    # Return only the host outputs.    return [out.host for out in outputs]
    
    def postprocess_the_outputs(h_outputs, shape_of_output):    h_outputs = h_outputs.reshape(*shape_of_output)    return h_outputs
    
    
    img_np_nchw = get_img_np_nchw(filename)img_np_nchw = img_np_nchw.astype(dtype=np.float32)
    # These two modes are dependent on hardwaresfp16_mode = Falseint8_mode = Falsetrt_engine_path = './model_fp16_{}_int8_{}.trt'.format(fp16_mode, int8_mode)# Build an engineengine = get_engine(max_batch_size, onnx_model_path, trt_engine_path, fp16_mode, int8_mode)# Create the context for this enginecontext = engine.create_execution_context()# Allocate buffers for input and outputinputs, outputs, bindings, stream = allocate_buffers(engine) # input, output: host # bindings
    # Do inferenceshape_of_output = (max_batch_size, 1000)# Load data to the bufferinputs[0].host = img_np_nchw.reshape(-1)
    # inputs[1].host = ... for multiple inputt1 = time.time()trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) # numpy datat2 = time.time()feat = postprocess_the_outputs(trt_outputs[0], shape_of_output)
    print('TensorRT ok')#将model改为自己的模型,此处为pytoch的resnet50,需联网下载model = torchvision.models.resnet50(pretrained=True).cuda()resnet_model = model.eval()
    input_for_torch = torch.from_numpy(img_np_nchw).cuda()t3 = time.time()feat_2= resnet_model(input_for_torch)t4 = time.time()feat_2 = feat_2.cpu().data.numpy()print('Pytorch ok!')
    
    mse = np.mean((feat - feat_2)**2)print("Inference time with the TensorRT engine: {}".format(t2-t1))print("Inference time with the PyTorch model: {}".format(t4-t3))print('MSE Error = {}'.format(mse))
    print('All completed!')

    报错:

    In node -1 (importModel): INVALID_VALUE: Assertion failed: !_importer_ctx.network()->hasImplicitBatchDimension() && "This version of the ONNX parser only supports TensorRT INetworkDefinitions with an explicit batch dimension. Please ensure the network was created using the EXPLICIT_BATCH NetworkDefinitionCreationFlag."
     
    

    解决:

        def build_engine(max_batch_size, save_engine):         EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)        with trt.Builder(TRT_LOGGER) as builder, \                builder.create_network(EXPLICIT_BATCH) as network, \                trt.OnnxParser(network, TRT_LOGGERas parser:

    报错:​​​​​​​

    Traceback (most recent call last):  line 126, in <listcomp>    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]pycuda._driver.LogicError: cuMemcpyHtoDAsync failed: invalid argument

    解决:​​​​​​​

    def get_img_np_nchw(filename):    image = cv2.imread(filename)    image_cv = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)    image_cv = cv2.resize(image_cv, (1920, 1080))

    输入的检测图像尺寸需要resize成model的input的size

    改为​​​​​​​

    def get_img_np_nchw(filename):    image = cv2.imread(filename)    image_cv = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)    image_cv = cv2.resize(image_cv, (544,544))

    报错​​​​​​​

    line 139, in postprocess_the_outputs    h_outputs = h_outputs.reshape(*shape_of_output)ValueError: cannot reshape array of size 5780 into shape (1,1000)

    解决:​​​​​​​

    #shape_of_output = (max_batch_size, 1000)#修改成自己模型ouput的大小shape_of_output = (1,20,17,17)
    

    二、PyTorch-ONNX 精度对齐工具

    精度对齐,是模型部署中重要的一个环节。在把深度学习框架模型转换成中间表示模型后,部署工程师们要做的第一件事就是精度对齐,确保模型的计算结果与之前相当。精度对齐时最常用的方法,就是使用测试集评估一遍中间表示模型,看看模型的评估指标(如准确度、相似度)是否下降。

    而在 PyTorch 到 ONNX 这条部署路线上,这种精度对齐方式有一些不便:一旦我们发现 PyTorch 模型和 ONNX 模型的评估指标有了出入,我们很难去追踪精度是在哪一个模块出了问题。这是因为 PyTorch 和 ONNX 模块总是难以对应。如下面的例子所示:

     

    假设我们现在有一个由很多卷积块 convs1, convs2... 组成的网络,我们想对齐 PyTorch 模型和 ONNX 模型的精度。第一步,我们想比较第一个卷积块的输出 x = self.convs1(x)。模块在PyTorch 模型中的输出可以很轻松地得到,可是,这个输出究竟对应 ONNX 模型里的哪一个输出呢?在小模型里,我们或许能够通过阅读 PyTorch 模型的源码,推断出每个 ONNX 模块与 PyTorch 模块的对应关系;但是,在大模型中,我们是难以建立 PyTorch 与 ONNX 的对应关系的。延庆川北小区45孙老师 收卖废品破烂垃圾炒股 废品孙 再回收

    在这篇教程中,我们就来利用自定义算子、子模型提取等工具,实现一个简单的 PyTorch-ONNX 精度对齐工具。

    设计思路

    为了把 PyTorch 和 ONNX 模块对应起来,我们可以使用一种储存了调试信息的自定义算子,如下图所示:

     

    我们可以定义一个叫做 Debug 的 ONNX 算子,它有一个属性调试名 name。而由于每一个 ONNX 算子节点又自带了输出张量的名称,这样一来,ONNX 节点的输出名和调试名绑定在了一起。我们可以顺着 PyTorch 里的调试名,找到对应 ONNX 里的输出,完成 PyTorch 和 ONNX 的对应。

    比如在上图的例子中,我们把第一个卷积块输出 x=self.convs1(x) 接入一个带有调试名 x_0 的调试算子。在最后生成的 ONNX 模型中,假设调试名 x_0 对应的输出张量叫做 a。知道了这一信息后,我们只需要先运行一遍 PyTorch 模型,记录第一个卷积块的输出;再运行一遍 ONNX 模型,用上篇教程中提到的截取 ONNX 中间结果的方法,记录中间张量 a 的值。这样,我们就可以对齐某 PyTorch 模块和它对应的 ONNX 模块的输出了。

    代码实现

    Debug 算子

    首先,我们需要实现之前提到的 Debug 算子:

    1. import torch 
    2.  
    3. class DebugOp(torch.autograd.Function): 
    4.     @staticmethod 
    5.     def forward(ctx, x, name): 
    6.         return x 
    7.  
    8.     @staticmethod 
    9.     def symbolic(g, x, name): 
    10.         return g.op("my::Debug", x, name_s=name) 
    11.  
    12. debug_apply = DebugOp.apply 

    Debug 算子的调用接口有两个参数:输入张量 x 和调试名 name。为了把这个算子“伪装”成一个普通的算子,使之能正常地参与推理、构建计算图的操作,我们还是需要正确定义对输入 x 进行操作的 forward 函数。而在表示 PyTorch 与 ONNX 映射规则的 symbolic 函数里,我们要定义一个带有调试名的 ONNX 算子,并把输入的 name 传给算子。

    由于 Debug 算子本身不表示任何计算,因此在 forward 函数中,直接把输入 x 返回即可。

    而 symbolic 函数定义了一个新算子 my::Debug:算子有一个输入 x,一个属性 name。我们直接把算子调用接口里的 xname 传入即可。

    这里需要补充介绍算子定义函数 g.op() 的一些规范。在g.op()中,算子的属性需要以 {attibute_name}_{type}=attibute_value 这样的格式传入。其中 {attibute_name} 为属性名,{type} 指定了算子属性的数据类型。比如说我们上面的算子属性写成 name_s,实际上是定义了一个字符串类型,名字叫做 name 的属性。除了表示字符串类型的 _s 外,还有表示 float 型的 _f,表示 tensor 型的 _t

    在完成算子的定义后,我们可以通过 debug_apply = DebugOp.apply 获取算子的调用接口。这样以后就可以通过 debug_apply(x, name) 来使用这个算子了。

    如果对 torch.autograd.Function 的用法不熟,欢迎回顾第四篇教程。

    Debugger 类

    接着,我们来实现精度对齐工具的核心——Debugger 类。这个类包含了实现精度对齐所需的所有操作。其定义如下:

     

    Debugger 类有三个成员变量:

    • torch_value 记录了运行 PyTorch 模型后每个调试张量的值。

    • onnx_value 记录了运行 ONNX 模型后每个调试张量的值。

    • output_debug_name 记录了把调试张量加入 ONNX 的输出后,每个输出张量的调试名。

    稍后我们会在类实现的代码里看到这些成员变量的具体用法。

    Debugger 类有以下方法:

    • debug 封装了之前编写好的 debug_apply。该方法需要在原 PyTorch 模型中调用,可以为导出的 ONNX 模型添加 Debug 算子节点,同时记录 PyTorch 调试张量值。

    • extract_debug_model 和 ONNX 的子模型提取函数的用法类似,可以把带调试节点的 ONNX 模型转化成一个可以输出调试张量的 ONNX 模型。

    • run_debug_model 会使用 ONNX Runtime 运行模型,得到 ONNX 调试张量值。

    • print_debug_result 会比较 PyTorch 和 ONNX 的调试张量值,输出比较的结果。

    这 4 个方法会依次被调用。下面我们来具体探究一下每个方法的实现。

    生成调试节点

    1. def debug(self, x, name)
    2.     self.torch_value[name] = x.detach().cpu().numpy() 
    3.     return debug_apply(x, name) 

    如前文所述,debug完成了两件事:记录 PyTorch 模型中调试张量的值、添加 Debug 节点。我们使用 self.torch_value[name] = x.detach().cpu().numpy() 把调试张量转成 numpy 格式并保存进 torch_value 词典里。之后,我们调用之前编写的 debug_apply 算子。

    提取调试模型

    1. def extract_debug_model(self, input_path, output_path): 
    2.     model = onnx.load(input_path) 
    3.     inputs = [input.name for input in model.graph.input
    4.     outputs = [] 
    5.  
    6.     for node in model.graph.node: 
    7.         if node.op_type == 'Debug'
    8.             # 记录调试张量名 
    9.             debug_name = node.attribute[0].s.decode('ASCII'
    10.             self.output_debug_name.append(debug_name) 
    11.  
    12.             # 添加输入 
    13.             output_name = node.output[0
    14.             outputs.append(output_name) 
    15.  
    16.             # 转换 Debug 节点为 Indentity 节点 
    17.             node.op_type = 'Identity' 
    18.             node.domain = '' 
    19.             del node.attribute[:] 
    20.  
    21.     e = onnx.utils.Extractor(model) 
    22.     extracted = e.extract_model(inputs, outputs) 
    23.     onnx.save(extracted, output_path) 

    在 PyTorch 模型中插入 debug 方法后,我们可以得到一个包含了若干 Debug 节点的 ONNX 模型。但是,这个 ONNX 模型不是我们最终拿来执行的模型。为了得到 Debug 节点的输出(即调试张量的值),我们需要做三项处理以提取出一个可运行的调试模型:

    1. 记录每个调试张量的调试名,为之后对齐 PyTorch、ONNX 调试张量值做准备。

    2. 把所有 Debug 节点的输出加入到整个模型的输出中,这样在运行模型后就能得到这些中间节点的输出了。

    3. 自定义的 Debug 节点在推理引擎中是没有实现的,为了让处理后的 ONNX 模型运行起来,需要把 Debug 节点转化成可运行的 Identity (恒等)节点。

    完成了这三项处理后,我们才能进行模型提取。下面,我们来看看模型提取和这几项处理是怎么实现的。
    首先,看一下和模型提取有关的代码:

    1. model = onnx.load(input_path) 
    2. inputs = [input.name for input in model.graph.input
    3. outputs = [] 
    4.  
    5. # 获取 outputs 
    6. ... 
    7.  
    8. # 调用提取模型 API 
    9. e = onnx.utils.Extractor(model) 
    10. extracted = e.extract_model(inputs, outputs) 
    11.  
    12. # 保存模型 
    13. onnx.save(extracted, output_path) 

    在提取模型时,我们要准备新模型的输入和输出。输入张量 inputs 还是保持原状,而输出张量 outputs 会在之后填入 Debug 节点的输出。获取完 outputs 后,我们调用提取模型的 API,得到处理过后的模型,并保存此模型。

    接着,看一下主处理逻辑:

    1. for node in model.graph.node: 
    2.     if node.op_type == 'Debug': 
    3.         ... 

    为了获取和 Debug 节点相关的信息,我们需要遍历 ONNX 模型的所有节点,找出那些类型为 Debug 的节点,并对这些节点执行操作。

    下面的代码实现了记录调试张量名:

    1. debug_name = node.attribute[0].s.decode('ASCII')
    2. self.output_debug_name.append(debug_name)

    这段代码的作用是:从节点的第一个属性(即 name )中取出调试名信息,并存入 output_debug_name 中。节点第一个属性的值可以通过 node.attribute[0] 获得。由于 name 是属性是字符串,这里要用 .s 获取属性的字符串值。又由于 ONNX 是以二进制的形式保存所有数据的,这里要用 .decode('ASCII') 把二进制字符串转成一个文本字符串。

    接下来的代码用于填写新模型输出 outputs

    1. output_name = node.output[0
    2. outputs.append(output_name) 

    node.output[0] 就是 Debug 节点的输出张量在 ONNX 里的名称。把这个名称加入新模型的输出后,只要运行新模型,就可以得到该输出张量的值了。

    最后这段代码用于更改 Debug 节点的类型:

    1. node.op_type = 'Identity' 
    2. node.domain = '' 
    3. del node.attribute[:] 

    为了消除 ONNX 不支持的 Debug 节点,一种比较简单的方式是直接把 Debug 节点修改成不执行任何操作的 Indentity 类型的节点。为了做这个转换,我们要先修改节点类型名 node.op_type 为Identity,再把节点的域(即命名空间)node.domain 修改成空,最后删除节点的所有属性,保证节点符合 ONNX 的规范。
    回忆一下,如果一个节点的 domain 为空,这个节点就会被当成一个 ONNX 原生算子节点。

    运行调试模型

    在生成调试节点时,我们已经顺便记录了 PyTorch 模型调试张量的值。下一步,我们要运行调试模型,记录 ONNX 模型调试张量的值。其实现如下:

    1. def run_debug_model(self, input, debug_model): 
    2.     sess = onnxruntime.InferenceSession(debug_model,  
    3.     providers=['CPUExecutionProvider']) 
    4.     onnx_outputs = sess.run(Noneinput
    5.  
    6.     for name, value in zip(self.output_debug_name, onnx_outputs): 
    7.         self.onnx_value[name] = value 

    在运行调试模型前,我们要给出模型输入、模型名这两个参数。根据这些参数,run_debug_model 会调用 ONNX Runtime 的 API,对 ONNX 模型进行推理。在得到了 ONNX 模型的输出后,我们要使用上一步得到的 output_debug_name 信息,填写 onnx_value,把 ONNX 的中间运算结果绑定到调试名上。完成了这些步骤之后,我们就有足够的信息做精度对齐了。

    输出调试信息

    1.  def print_debug_result(self): 
    2.     for name in self.torch_value.keys(): 
    3.         if name in self.onnx_value: 
    4.             mse = np.mean(self.torch_value[name] - self.onnx_value[name])**2
    5.             print(f"{name} MSE: {mse}"

    最后,我们同时遍历 self.torch_value 和 self.onnx_value 这两个词典,比较同一个张量在 PyTorch 模型和 ONNX 模型里的输出。在循环体中,我们只需要使用 self.torch_value[name] 和 self.onnx_value[name] 就可以访问同一个张量在 PyTorch 里的值和在 ONNX 里的值。作为示例,这里我们可以计算二者的均方误差 mse,以此为精度对齐的依据。

    整理一下,整个工具库的代码如下:

    1. import torch 
    2. import onnx 
    3. import onnxruntime 
    4. import numpy as np 
    5.  
    6. class DebugOp(torch.autograd.Function): 
    7.     @staticmethod 
    8.     def forward(ctx, x, name): 
    9.         return x 
    10.  
    11.     @staticmethod 
    12.     def symbolic(g, x, name): 
    13.         return g.op("my::Debug", x, name_s=name) 
    14.  
    15. debug_apply = DebugOp.apply 
    16.  
    17. class Debugger(): 
    18.     def __init__(self): 
    19.         super().__init__() 
    20.         self.torch_value = dict() 
    21.         self.onnx_value = dict() 
    22.         self.output_debug_name = [] 
    23.  
    24.     def debug(self, x, name): 
    25.         self.torch_value[name] = x.detach().cpu().numpy() 
    26.         return debug_apply(x, name) 
    27.  
    28.     def extract_debug_model(self, input_path, output_path): 
    29.         model = onnx.load(input_path) 
    30.         inputs = [input.name for input in model.graph.input
    31.         outputs = [] 
    32.  
    33.         for node in model.graph.node: 
    34.             if node.op_type == 'Debug'
    35.                 debug_name = node.attribute[0].s.decode('ASCII'
    36.                 self.output_debug_name.append(debug_name) 
    37.  
    38.                 output_name = node.output[0
    39.                 outputs.append(output_name) 
    40.  
    41.                 node.op_type = 'Identity' 
    42.                 node.domain = '' 
    43.                 del node.attribute[:] 
    44.         e = onnx.utils.Extractor(model) 
    45.         extracted = e.extract_model(inputs, outputs) 
    46.         onnx.save(extracted, output_path) 
    47.  
    48.     def run_debug_model(self, input, debug_model): 
    49.         sess = onnxruntime.InferenceSession(debug_model,  
    50.         providers=['CPUExecutionProvider']) 
    51.  
    52.         onnx_outputs = sess.run(Noneinput
    53.         for name, value in zip(self.output_debug_name, onnx_outputs): 
    54.             self.onnx_value[name] = value 
    55.  
    56.     def print_debug_result(self): 
    57.         for name in self.torch_value.keys(): 
    58.             if name in self.onnx_value: 
    59.                 mse = np.mean(self.torch_value[name] - self.onnx_value[name])**2
    60.                 print(f"{name} MSE: {mse}"

    使用方法

    实现了精度对齐工具后,我们来看看该怎么把这个工具用起来。

    现在,假设我们得到了一个这样的模型:

    1. class Model(torch.nn.Module):
    2. def __init__(self):
    3. super().__init__()
    4. self.convs1 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1),
    5. torch.nn.Conv2d(3, 3, 3, 1, 1),
    6. torch.nn.Conv2d(3, 3, 3, 1, 1))
    7. self.convs2 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1),
    8. torch.nn.Conv2d(3, 3, 3, 1, 1))
    9. self.convs3 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1),
    10. torch.nn.Conv2d(3, 3, 3, 1, 1))
    11. self.convs4 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3, 1, 1),
    12. torch.nn.Conv2d(3, 3, 3, 1, 1),
    13. torch.nn.Conv2d(3, 3, 3, 1, 1))
    14. def forward(self, x):
    15. x = self.convs1(x)
    16. x = self.convs2(x)
    17. x = self.convs3(x)
    18. x = self.convs4(x)
    19. return x
    20. torch_model = Model()

    没错!这就是本文开头展示的那个全卷积网络。现在我们想对齐 convs1 至 convs4 这每一个卷积块的输出精度,该怎么使用之前写好的精度对齐工具呢?

    首先,我们生成管理类 Debugger 的一个实例:

    debugger = Debugger() 
    

    之后,我们要设法把 Debug 节点插入原模型:

    1. from types import MethodType 
    2.  
    3. def new_forward(self, x): 
    4.     x = self.convs1(x) 
    5.     x = debugger.debug(x, 'x_0'
    6.     x = self.convs2(x) 
    7.     x = debugger.debug(x, 'x_1'
    8.     x = self.convs3(x) 
    9.     x = debugger.debug(x, 'x_2'
    10.     x = self.convs4(x) 
    11.     x = debugger.debug(x, 'x_3'
    12.     return x 
    13.  
    14. torch_model.forward = MethodType(new_forward, torch_model) 

    我们可以为原模型新写一个 forward 函数。在这个新的函数函数中,我们可以通过 debugger.debug 把每一个输出张量标记起来,并各取一个不重复的调试名。

    有了 new_forward 函数,我们需要使用 MethodType 这个 Python API 把这个函数变成模型实例 torch_model 的一个成员方法,确保 torch_model 的 forward 函数能够被正确替换。

    实现了”狸猫换太子“般巧妙的操作后,我们就可以使用 PyTorch API 导出一个带有 Debug 节点的 ONNX 模型了:

    1. dummy_input = torch.randn(131010
    2. torch.onnx.export(torch_model, dummy_input, 'before_debug.onnx', input_names=['input']) 

    由于 torch.onnx.export 模型使用的是跟踪法,模型的 forward 函数会被执行一次, debugger.debug 操作可以把 PyTorch 模型的调试张量输出记录在 debugger.torch_value 里。

    这个 before_debug.onnx 模型的部分可视化结果如下:

     

    接下来,我们替换掉所有 Debug 节点,并记录每个 Debug 输出张量的 ONNX 名与调试名的对应关系:

    debugger.extract_debug_model('before_debug.onnx''after_debug.onnx'

    这步操作得到的 after_debug.onnx 模型的部分可视化结果如下:

     

    我们可以使用下面的代码运行这个模型:

    debugger.run_debug_model({'input':dummy_input.numpy()}, 'after_debug.onnx')
    

    这样,ONNX 模型的调试张量输出会记录在 debugger.onnx_value 里。

    总算,一切准备工作结束了。我们可以轻轻松松地用一行代码输出精度对齐的结果:

    debugger.print_debug_result() 
    

    这个函数大致会输出以下内容:

    1. x_0 MSE: 8.465450562766819e-16  
    2. x_1 MSE: 1.4122021817221354e-16  
    3. x_2 MSE: 6.501743508551734e-17  
    4. x_3 MSE: 1.7635199492054931e-16 

    这份输出表明,在这一轮精度对齐测试中,所有模块的精度误差都很小。我们几乎可以认为,ONNX 模型的运行结果等价于 PyTorch 模型的运行结果。whaosoft aiot http://143ai.com

    如果有某些模块的误差比较大,我们可以深入子模块,去加更多的 debug 节点,看看是哪一步、哪一个算子出现了问题。

    总结

    基于前几篇教程中的知识,本文我们介绍了一个与 PyTorch 转 ONNX 相关的实战项目。

    其实,这篇教程提到的精度对齐工具还有很大的优化空间,比如:

    • 现在测试输入dummy_input不太方便管理,能不能优化它的管理逻辑?

    • 现在我们默认是用 ONNX Runtime 执行 ONNX 模型,可否拓展到其他后端?

    • 能不能把无效的 Identity 节点从调试模型中删除?

  • 相关阅读:
    浅谈Rust内存管理
    2022年最新西藏机动车签字授权人模拟考试及答案
    隐私计算 FATE - 多分类神经网络算法测试
    CSS系列之文本换行
    「津津乐道播客」#389 科技乱炖:技术播客月,我们一起聊聊技术与开源
    35m预应力简支梁桥毕业设计 课程设计-桥梁工程(计算书、8张CAD图)
    Codeforces Round #801 (Div. 2) and EPIC Institute of Technology Round(C,D题解)
    【spring】bean的生命周期
    科技前沿:Web3与物联网的智能连接
    ES6 入门教程 6 正则的扩展 6.10 Unicode 属性类 & 6.11 v 修饰符:Unicode 属性类的运算 & 6.12 具名组匹配
  • 原文地址:https://blog.csdn.net/qq_29788741/article/details/126018375