onnx代码解读

🕗 发布于 2024-10-12 10:32 自然语言处理

一、定义

torch.jit.trace 相关代码解读
onnx 内部实现
3 查看是否为aten 算子
aten 算子实现
torch.autograd.Functions 算子实现
自定义算子实现
查找未实现的节点
一次性发现所有的未实现 aten 算子

二、实现

torch.jit.trace 相关代码解读
1. torch.jit.script() : 将其转换为可运行的脚本。转换后的脚本可以像普通的 Python 函数一样调用，也可以保存到磁盘并在没有 PyTorch 依赖的环境中执行。
2. torch.jit.trace : 跟踪了给定输入张量的执行路径，因此在使用转换后的模块对象进行推理时，输入张量的维度和数据类型必须与跟踪时使用的相同。

3 查看是否为aten 算子

import torch

print(
    torch.jit.trace(
        torch.nn.ELU(), # module
        torch.ones(1)   # example input
    ).graph
)

算子追踪，在这里插入图片描述
3. aten 算子实现
　　１．查看ｔｏｒｃｈ　接口定义　　　　torch/nn/functional.pyi
　　２．查看ｏｎｎｘ　算子命名　　　　https://github.com/onnx/onnx/blob/main/docs/Operators.md
　　3.　查看注册函数书写　　　symbolic_opset9.py

import torch
import torch.onnx.verification
torch.manual_seed(0)
opset_version = 15
# Define a custom symbolic function for aten::relu.
# The custom symbolic function is incorrect, which will result in mismatches.

#def relu(input: Tensor) -> Tensor: ...   查看接口定义，
def correct_relu_symbolic_function(g, input):
    return g.op("Relu", input)             #查看onnx 实现

torch.onnx.register_custom_op_symbolic(     #注册
    "aten::relu",
    correct_relu_symbolic_function,
    opset_version=opset_version,
)

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(3, 4),
            torch.nn.ReLU(),
            torch.nn.Linear(4, 5),
            torch.nn.ReLU(),
            torch.nn.Linear(5, 6),
        )
    def forward(self, x):
        return self.layers(x)

graph_info = torch.onnx.verification.find_mismatch(
    Model(),
    (torch.randn(2, 3),),
    opset_version=opset_version,
)

torch.autograd.Functions 算子实现
如果算子是torch.autograd.Functions 的子模块，可以使用该方法实现。

import torch

class MyRelu(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input: torch.Tensor) -> torch.Tensor:
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    @staticmethod
    def symbolic(g: torch.Graph, input: torch.Value) -> torch.Value:
        return g.op("Clip", input, g.op("Constant", value_t=torch.tensor(0, dtype=torch.float)))


import torch
import torch.onnx.verification
torch.manual_seed(0)
opset_version = 15

myrelu = MyRelu.apply　　　　　　　　＃核心
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(3, 4),
            torch.nn.Linear(4, 5),
            torch.nn.Linear(5, 6),
        )
    def forward(self, x):
        return myrelu(self.layers(x))

graph_info = torch.onnx.verification.find_mismatch(
    Model(),
    (torch.randn(2, 3),),
    opset_version=opset_version,
)

自定义算子实现
1.　onnx 算子实现
1. 自定义c++ 算子 +Extending TorchScript with Custom C++ Operators 实现
查找未实现的节点

import torch
import torch.onnx.verification
torch.manual_seed(0)
opset_version = 15
# Define a custom symbolic function for aten::relu.
# The custom symbolic function is incorrect, which will result in mismatches.   注册函数错误，导致find_mismatch 算子
def incorrect_relu_symbolic_function(g, self):
    return self
torch.onnx.register_custom_op_symbolic(
    "aten::relu",
    incorrect_relu_symbolic_function,
    opset_version=opset_version,
)
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(3, 4),
            torch.nn.ReLU(),
            torch.nn.Linear(4, 5),
            torch.nn.ReLU(),
            torch.nn.Linear(5, 6),
        )
    def forward(self, x):
        return self.layers(x)
graph_info = torch.onnx.verification.find_mismatch(
    Model(),
    (torch.randn(2, 3),),
    opset_version=opset_version,
)

#===================== Mismatch info for graph partition : ======================
================================ Mismatch error ================================
Tensor-likes are not close!
Mismatched elements: 12 / 12 (100.0%)
Greatest absolute difference: 0.2328854203224182 at index (1, 2) (up to 1e-07 allowed)
Greatest relative difference: 0.699536174352349 at index (1, 3) (up to 0.001 allowed)
==================================== Tree: =====================================
5 X   __2 X    __1 \u2713
id:  |  id: 0 |  id: 00
     |        |
     |        |__1 X (aten::relu)
     |           id: 01
     |
     |__3 X    __1 \u2713
        id: 1 |  id: 10
              |
              |__2 X     __1 X (aten::relu)
                 id: 11 |  id: 110
                        |
                        |__1 \u2713
                           id: 111
=========================== Mismatch leaf subgraphs: ===========================
['01', '110']
============================= Mismatch node kinds: =============================
{'aten::relu': 2}

修改后：
aten 算子实现

import torch
import torch.onnx.verification
torch.manual_seed(0)
opset_version = 15
# Define a custom symbolic function for aten::relu.
# The custom symbolic function is incorrect, which will result in mismatches.

#def relu(input: Tensor) -> Tensor: ...   查看接口定义，
def correct_relu_symbolic_function(g, input):
    return g.op("Relu", input)             #查看onnx 实现

torch.onnx.register_custom_op_symbolic(     #注册
    "aten::relu",
    correct_relu_symbolic_function,
    opset_version=opset_version,
)

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(3, 4),
            torch.nn.ReLU(),
            torch.nn.Linear(4, 5),
            torch.nn.ReLU(),
            torch.nn.Linear(5, 6),
        )
    def forward(self, x):
        return self.layers(x)

graph_info = torch.onnx.verification.find_mismatch(
    Model(),
    (torch.randn(2, 3),),
    opset_version=opset_version,
)

方式二、
c++ 自定义算子


import torch
import torch.onnx.verification
torch.manual_seed(0)
opset_version = 15


from torch.onnx import register_custom_op_symbolic        # 为 TorchScript 算子补充注册符号函数
from torch.onnx.symbolic_helper import parse_args
# '''
# 装饰器 @parse_args 了。简单来说，TorchScript 算子的符号函数要求标注出每一个输入参数的类型。比如"v"表示 Torch 库里的 value 类型，
# 一般用于标注张量，而"i"表示 int 类型，"f"表示 float 类型，"none"表示该参数为空。具体的类型含义可以在 torch.onnx.symbolic_helper.py
# '''
@parse_args("v", "v")
def correct_relu_symbolic_function(g,input):
    return g.op("Relu", input)


torch.onnx.register_custom_op_symbolic(     #注册
    "aten::relu",
    correct_relu_symbolic_function,
    opset_version=opset_version,
)

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(3, 4),
            torch.nn.ReLU(),
            torch.nn.Linear(4, 5),
            torch.nn.ReLU(),
            torch.nn.Linear(5, 6),
        )
    def forward(self, x):
        return self.layers(x)

graph_info = torch.onnx.verification.find_mismatch(
    Model(),
    (torch.randn(2, 3),),
    opset_version=opset_version,
)

一次性发现所有的未实现 aten 算子

import torch
import torch.onnx.verification
torch.manual_seed(0)
opset_version = 15

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(3, 4),
            torch.nn.ReLU(),
            torch.nn.Linear(4, 5),
            torch.nn.ReLU(),
            torch.nn.Linear(5, 6),
        )
    def forward(self, x):
        return self.layers(x)


torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(
    Model(), (torch.randn(2, 3),), opset_version=opset_version
)

print(set(unconvertible_ops))

原文地址：https://blog.csdn.net/weixin_40777649/article/details/142858767

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！

上一篇：Elasticsearch实战应用
下一篇：YOLO11改进|注意力机制篇|引入轴向注意力Axial Attention

下一代安全：融合网络和物理策略以实现最佳保护
网络安全和物理安全的融合不仅是一种趋势，而且是现代威胁形势下的必需品。
阅读更多2024-10-15
开发规范 - mac系统1小时装机极速装机开发环境
jdk不用单配配置，在idea中，选择一个语言环境（jdk8/jdk11/jdk17…）,然后默认下载jdk这样，一个hello word就能跑起来了。
阅读更多2024-10-15
CsQuery操作详解
CsQuery 是一个强大且易于使用的 C# 库，适合需要类似 jQuery 的直观 API 来解析和操作 HTML 文档的场景。通过详细的示例，你可以看到如何利用 CsQuery 执行各种常见的 H
阅读更多2024-10-15
大厂面试真题-说一说Kafka的消息积压
Kafka消息积压是指Kafka中消息未能被消费者及时消费而在队列中累积的现象。
阅读更多2024-10-15
【HarmonyOS】HMRouter使用详解（四）路由拦截
可以对指定或全局路由跳转时添加拦截器，作用是可以实现在页面切换前做判断是否有进入当前页面的权限。这篇文章将实现登录的全局路由拦截样式。
阅读更多2024-10-15
使用开源的 Vue 移动端表单设计器创建表单
FcDesigner Vant 版是一款开源的基于 Vue3.0 的移动端低代码可视化表单设计器工具，通过数据驱动表单渲染。可以通过拖拽的方式快速创建表单，提高开发者对表单的开发效率，节省开发者的时间
阅读更多2024-10-15
LEETCODE 49场周赛第K大完美二叉子树的大小
第K大完美二叉子树
阅读更多2024-10-15
入门网络安全工程师要学习哪些内容
大家都知道网络安全行业很火，这个行业因为国家政策趋势正在大力发展，大有可为!但很多人对还是不了解，不知道需要学什么?知了堂小编总结出以下要点。是一个概称，学习的东西很多，具体学什么看自己以后的职业定位
阅读更多2024-10-15
git rebase
是一个强大的工具，允许你灵活地管理和修改提交历史。它通常用于整理提交历史、合并冗余提交、修改提交信息等操作。使用时需要小心，特别是在公共分支上，因为它会重写历史，可能会影响其他开发者。参考：官网：ht
阅读更多2024-10-15
一个好的维权小程序应该是什么样的？
小程序的直接目的都是为了服务用户，一切应该以用户为主，所以一个好的维权小程序不仅能够为用户提供便捷、高效的服务，还能在维护社会公平正义方面发挥积极作用。，对于所受到的投诉都进行反馈和处理，用户能够看到
阅读更多2024-10-15

onnx代码解读

一、定义

二、实现

相关文章