使用Pytorch导出自定义ONNX算子
在实际部署模型时有时可能会遇到想用的算子无法导出onnx,但实际部署的框架是支持该算子的。此时可以通过自定义onnx算子的方式导出onnx模型(注:自定义onnx算子导出onnx模型后是无法使用onnxruntime推理的)。下面给出个具体应用中的示例:需要导出pytorch的affine_grid算子,但在pytorch的2.0.1版本中又无法正常导出该算子,故可通过如下自定义算子代码导出。
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function
from torch.onnx import OperatorExportTypesclass CustomAffineGrid(Function):@staticmethoddef forward(ctx, theta: torch.Tensor, size: torch.Tensor):grid = F.affine_grid(theta=theta, size=size.cpu().tolist())return grid@staticmethoddef symbolic(g: torch.Graph, theta: torch.Tensor, size: torch.Tensor):return g.op("AffineGrid", theta, size)class MyModel(nn.Module):def __init__(self) -> None:super().__init__()def forward(self, x: torch.Tensor, theta: torch.Tensor, size: torch.Tensor):grid = CustomAffineGrid.apply(theta, size)x = F.grid_sample(x, grid=grid, mode="bilinear", padding_mode="zeros")return xdef main():with torch.inference_mode():custum_model = MyModel()x = torch.randn(1, 3, 224, 224)theta = torch.randn(1, 2, 3)size = torch.as_tensor([1, 3, 512, 512])torch.onnx.export(model=custum_model,args=(x, theta, size),f="custom.onnx",input_names=["input0_x", "input1_theta", "input2_size"],output_names=["output"],dynamic_axes={"input0_x": {2: "h0", 3: "w0"},"output": {2: "h1", 3: "w1"}},opset_version=16,operator_export_type=OperatorExportTypes.ONNX_FALLTHROUGH)if __name__ == '__main__':main()
在上面代码中,通过继承torch.autograd.Function父类的方式实现导出自定义算子,继承该父类后需要用户自己实现forward以及symbolic两个静态方法,其中forward方法是在pytorch正常推理时调用的函数,而symbolic方法是在导出onnx时调用的函数。对于forward方法需要按照正常的pytorch语法来实现,其中第一个参数必须是ctx但对于当前导出onnx场景可以不用管它,后面的参数是实际自己传入的参数。对于symbolic方法的第一个必须是g,后面的参数任为实际自己传入的参数,然后通过g.op方法指定具体导出自定义算子的名称,以及输入的参数(注:上面示例中传入的都是Tensor所以可以直接传入,对与非Tensor的参数可见下面一个示例)。最后在使用时直接调用自己实现类的apply方法即可。使用netron打开自己导出的onnx文件,可以看到如下所示网络结构。

有时按照使用的推理框架导出自定义算子时还需要设置一些参数(非Tensor)那么可以参考如下示例,例如要导出int型的参数k那么可以通过传入k_i来指定,要导出float型的参数scale那么可以通过传入scale_f来指定,要导出string型的参数clockwise那么可以通过传入clockwise_s来指定:
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Function
from torch.onnx import OperatorExportTypesclass CustomRot90AndScale(Function):@staticmethoddef forward(ctx, x: torch.Tensor):x = torch.rot90(x, k=1, dims=(3, 2)) # clockwise 90x *= 1.2return x@staticmethoddef symbolic(g: torch.Graph, x: torch.Tensor):return g.op("Rot90AndScale", x, k_i=1, scale_f=1.2, clockwise_s="yes")class MyModel(nn.Module):def __init__(self) -> None:super().__init__()def forward(self, x: torch.Tensor):return CustomRot90AndScale.apply(x)def main():with torch.inference_mode():custum_model = MyModel()x = torch.randn(1, 3, 224, 224)torch.onnx.export(model=custum_model,args=(x,),f="custom_rot90.onnx",input_names=["input"],output_names=["output"],dynamic_axes={"input": {2: "h0", 3: "w0"},"output": {2: "w0", 3: "h0"}},opset_version=16,operator_export_type=OperatorExportTypes.ONNX_FALLTHROUGH)if __name__ == '__main__':main()
使用netron打开自己导出的onnx文件,可以看到如下所示信息。

相关文章:
使用Pytorch导出自定义ONNX算子
在实际部署模型时有时可能会遇到想用的算子无法导出onnx,但实际部署的框架是支持该算子的。此时可以通过自定义onnx算子的方式导出onnx模型(注:自定义onnx算子导出onnx模型后是无法使用onnxruntime推理的)。下面给出个具体应用中的…...
unity-urp:视野雾
问题背景 恐怖游戏在黑夜或者某些场景下,需要用雾或者黑暗遮盖视野,搭建游戏氛围 效果 场景中,雾会遮挡场景和怪物,但是在玩家视野内雾会消散,距离玩家越近雾越薄。 当前是第三人称视角,但是可以轻松的…...
Spring Cloud Gateway介绍及入门配置
Spring Cloud Gateway介绍及入门配置 概述: Gateway是在Spring生态系统之上构建的API网关服务,基于Spring6,Spring Boot 3和Project Reactor等技术。它旨在为微服务架构提供一种简单有效的统一的 API 路由管理方式,并为它们提供…...
Thingsboard本地源码部署教程
本章将介绍ThingsBoard的本地环境搭建,以及源码的编译安装。本机环境:jdk11、maven 3.6.2、node v12.18.2、idea 2023.1、redis 6.2 环境安装 开发环境要求: Jdk 11 版本 ;Postgresql 9 以上;Maven 3.6 以上…...
【MySQL 系列】MySQL 起步篇
MySQL 是一个开放源代码的、免费的关系型数据库管理系统。在 Web 开发领域,MySQL 是最流行、使用最广泛的关系数据库。MySql 分为社区版和商业版,社区版完全免费,并且几乎能满足全部的使用场景。由于 MySQL 是开源的,我们还可以根…...
C++的成员初始化列表
C的成员构造函数初始化列表:构造函数中初始化类成员的一种方式,当我们编写一个类并向该类添加成员时,通常需要某种方式对这些成员变量进行初始化。 建议应该在所有地方使用成员初始化列表进行初始化 成员初始化的方法 方法一: …...
为什么TikTok视频0播放?账号权重提高要重视
许多TikTok账号运营者都会遇到一个难题,那就是视频要么播放量很低,要么0播放!不管内容做的多好,最好都是竹篮打水一场空!其实你可能忽略了一个问题,那就是账号权重。下面好好跟大家讲讲这个东西!…...
element---tree树形结构(返回的数据与官方的不一样)
项目中要用到属性结构数据,后端返回的数据不是官方默认的数据结构: <el-tree:data"treeData":filter-node-method"filterNode":props"defaultProps"node-click"handleNodeClick"></el-tree>这是文档…...
Spring Boot工程集成验证码生成与验证功能教程
🌟 前言 欢迎来到我的技术小宇宙!🌌 这里不仅是我记录技术点滴的后花园,也是我分享学习心得和项目经验的乐园。📚 无论你是技术小白还是资深大牛,这里总有一些内容能触动你的好奇心。🔍 &#x…...
Bert Encoder和Transformer Encoder有什么不同
前言:本篇文章主要从代码实现角度研究 Bert Encoder和Transformer Encoder 有什么不同?应该可以帮助你: 深入了解Bert Encoder 的结构实现深入了解Transformer Encoder的结构实现 本篇文章不涉及对注意力机制实现的代码研究。 注:…...
外汇天眼:频繁交钱却无法出金,只因误入假冒HFM惨成冤大头!
在外汇市场上这么久了,天眼君总结出了一个不争的事实,但凡是不给出金或者以各种理由拒绝出金的平台一定有问题!想必不管是在外汇天眼还是其他地方,大家总是能看到一些外汇交易者投诉自己向平台申请出金需要缴纳各种费用࿰…...
Linux-信号3_sigaction、volatile与SIGCHLD
文章目录 前言一、sigaction__sighandler_t sa_handler;__sigset_t sa_mask; 二、volatile关键字三、SIGCHLD方法一方法二 前言 本章内容主要对之前的内容做一些补充。 一、sigaction #include <signal.h> int sigaction(int signum, const struct sigaction *act,struc…...
STM32 | STM32时钟分析、GPIO分析、寄存器地址查找、LED灯开发(第二天)
STM32 第二天 一、 STM32时钟分析 寄存器:寄存器的功能是存储二进制代码,它是由具有存储功能的触发器组合起来构成的。一个触发器可以存储1位二进制代码,故存放n位二进制代码的寄存器,需用n个触发器来构成 在计算机领域&#x…...
Python常用语法汇总(一):字符串、列表、字典操作
1. 字符串处理 print(message.title()) #首字母大写print(message.uper()) #全部大写print(message.lower()) #全部小写full_name "lin" "hai" #合并字符串print("Hello, " full_name.title() "!")print("John Q. %s10s&qu…...
Token的奥秘--一起学习吧之token
Token,在计算机科学中,是一个用于表示数据或一段数据的单位。它通常用于加密、身份验证、令牌化等场景,以确保数据的安全性和完整性。在编程语言中,Token通常是指代一段代码或数据的最小单元,例如一个变量、一个操作符…...
FlinkCDC快速搭建实现数据监控
引入依赖 <project xmlns"http://maven.apache.org/POM/4.0.0" xmlns:xsi"http://www.w3.org/2001/XMLSchema-instance"xsi:schemaLocation"http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd"><modelV…...
应急布控球远程视频监控方案:视频监控平台EasyCVR+4G/5G应急布控球
随着科技的不断发展,应急布控球远程视频监控方案在公共安全、交通管理、城市管理等领域的应用越来越广泛。这种方案通过在现场部署应急布控球,实现对特定区域的实时监控,有助于及时发现问题、快速响应,提高管理效率。 智慧安防视…...
3.6 C语言和汇编语言混合编程 “每日读书”
在一些嵌入式场合,我们经常看到C程序和汇编程序相互调用,混合编程,如在ARM启动代码中,系统上电首先运行的是汇编代码,等初始化好内存堆栈环境之后,才会跳到C程序中执行,对嵌入式软件进行优化时&…...
利用“定时执行专家”循环执行BAT、VBS、Python脚本——含参数指定功能
目录 一、软件概述 二、VBS脚本执行设置 三、触发器设置 四、功能亮点 五、总结 在自动化办公和日常计算机任务管理中,定时执行脚本是一项非常重要的功能。今天,我将为大家带来一款名为“定时执行专家”的软件的评测,特别是其定时执行VB…...
【算法集训】基础算法:模拟
一、基本理解 顾名思义,就是题目要求做什么,代码中就跟着做就可以。 二、题目练习 1252. 奇数值单元格的数目 根据题目要求列出如下代码。需要注意填充列和行的时候注意下标。 int oddCells(int m, int n, int** indices, int indicesSize, int* in…...
[2025CVPR]DeepVideo-R1:基于难度感知回归GRPO的视频强化微调框架详解
突破视频大语言模型推理瓶颈,在多个视频基准上实现SOTA性能 一、核心问题与创新亮点 1.1 GRPO在视频任务中的两大挑战 安全措施依赖问题 GRPO使用min和clip函数限制策略更新幅度,导致: 梯度抑制:当新旧策略差异过大时梯度消失收敛困难:策略无法充分优化# 传统GRPO的梯…...
Android Wi-Fi 连接失败日志分析
1. Android wifi 关键日志总结 (1) Wi-Fi 断开 (CTRL-EVENT-DISCONNECTED reason3) 日志相关部分: 06-05 10:48:40.987 943 943 I wpa_supplicant: wlan0: CTRL-EVENT-DISCONNECTED bssid44:9b:c1:57:a8:90 reason3 locally_generated1解析: CTR…...
微信小程序之bind和catch
这两个呢,都是绑定事件用的,具体使用有些小区别。 官方文档: 事件冒泡处理不同 bind:绑定的事件会向上冒泡,即触发当前组件的事件后,还会继续触发父组件的相同事件。例如,有一个子视图绑定了b…...
工业安全零事故的智能守护者:一体化AI智能安防平台
前言: 通过AI视觉技术,为船厂提供全面的安全监控解决方案,涵盖交通违规检测、起重机轨道安全、非法入侵检测、盗窃防范、安全规范执行监控等多个方面,能够实现对应负责人反馈机制,并最终实现数据的统计报表。提升船厂…...
从WWDC看苹果产品发展的规律
WWDC 是苹果公司一年一度面向全球开发者的盛会,其主题演讲展现了苹果在产品设计、技术路线、用户体验和生态系统构建上的核心理念与演进脉络。我们借助 ChatGPT Deep Research 工具,对过去十年 WWDC 主题演讲内容进行了系统化分析,形成了这份…...
将对透视变换后的图像使用Otsu进行阈值化,来分离黑色和白色像素。这句话中的Otsu是什么意思?
Otsu 是一种自动阈值化方法,用于将图像分割为前景和背景。它通过最小化图像的类内方差或等价地最大化类间方差来选择最佳阈值。这种方法特别适用于图像的二值化处理,能够自动确定一个阈值,将图像中的像素分为黑色和白色两类。 Otsu 方法的原…...
《通信之道——从微积分到 5G》读书总结
第1章 绪 论 1.1 这是一本什么样的书 通信技术,说到底就是数学。 那些最基础、最本质的部分。 1.2 什么是通信 通信 发送方 接收方 承载信息的信号 解调出其中承载的信息 信息在发送方那里被加工成信号(调制) 把信息从信号中抽取出来&am…...
Module Federation 和 Native Federation 的比较
前言 Module Federation 是 Webpack 5 引入的微前端架构方案,允许不同独立构建的应用在运行时动态共享模块。 Native Federation 是 Angular 官方基于 Module Federation 理念实现的专为 Angular 优化的微前端方案。 概念解析 Module Federation (模块联邦) Modul…...
【Zephyr 系列 10】实战项目:打造一个蓝牙传感器终端 + 网关系统(完整架构与全栈实现)
🧠关键词:Zephyr、BLE、终端、网关、广播、连接、传感器、数据采集、低功耗、系统集成 📌目标读者:希望基于 Zephyr 构建 BLE 系统架构、实现终端与网关协作、具备产品交付能力的开发者 📊篇幅字数:约 5200 字 ✨ 项目总览 在物联网实际项目中,**“终端 + 网关”**是…...
关于 WASM:1. WASM 基础原理
一、WASM 简介 1.1 WebAssembly 是什么? WebAssembly(WASM) 是一种能在现代浏览器中高效运行的二进制指令格式,它不是传统的编程语言,而是一种 低级字节码格式,可由高级语言(如 C、C、Rust&am…...
