解决低版本pytorch和onnx组合时torch.atan2()不被onnx支持的问题

解决这个问题,最简单的当然是升级pytorch和onnx到比较高的版本,例如有人验证过的组合: pytorch=2.1.1+cu118, onnxruntime=1.16.3

但是因为你的模型或cuda环境等约束,不能安装这么高的版本的pytorch和onnx组合时(例如我的环境是pytorch1.12,onnxruntime=1.19.2,即使onnxruntime版本比较高但是Pytorch的版本底也照样报这个错: tan2 to ONNX opset version 16 is not supported),那就只能考虑自己基于torch.atan()实现torch.atan2()的功能了。

torch.atan()因为不能区分坐标落在哪个像限,如果直接用来计算物体的朝向角的话,因为可能这个缺陷导致计算出来的物体的朝向是完全相反的方向,或者朝向沿着x轴做对称翻转了,例如

torch.atan(-1/10)和torch.atan(1/(-10))是没有区别的,atan2()就是为了解决这种问题的,atan2()的实现原理大致如下图所示:

有人基于atan()对atan2()做了如下近似实现:

复制代码
torch.atan((rot_sine / (rot_cosine + 1e-8)).sigmoid())                
+ ((1 - torch.sign(rot_cosine)) / 2) * torch.sign(rot_sine) * torch.pi

显然后半部分根据x和y的正负做加/减torch.pi是正确的,但是前半部分对(y/x)做sigmoid()运算把值一律转到不带符号的(0,1)之间在有的情况下是有一定误差的。

网上没找到其他更好的实现,于是我基于上面图中的计算规则做了如下实现:

复制代码
            rot_sine = bboxes[..., 6:7]
            rot_cosine = bboxes[..., 7:8]
           
            idx = torch.where(rot_cosine == 0)
            rot_cosine[idx] = 1e-6

            rot = torch.atan(rot_sine / rot_cosine)

            mask_yp = (rot_sine >= 0) & (rot_cosine < 0)
            mask_yn = (rot_sine < 0) & (rot_cosine < 0)
            idx_yp = torch.where(mask_yp)
            idx_yn = torch.where(mask_yn)
            rot[idx_yp] += torch.pi
            rot[idx_yn] -= torch.pi

用数据测试发现上述计算步骤计算出的结果和torch.atan2()计算出来是一致的,仅当x==0(或者说上面的rot_cosine==0)时,用小量1e-6代替0,计算出的角度和正负torch.pi/2可能有点很细微差异而已,这基本不影响物体朝向的正确性。

将上述实现封装成函数替代调用处的torch.atan2(),导出onnx就可以顺利成功了。

相关推荐
工程师老罗25 分钟前
基于Pytorch的YOLOv1 的网络结构代码
人工智能·pytorch·yolo
JarryStudy1 小时前
HCCL与PyTorch集成 hccl_comm.cpp DDP后端注册全流程
人工智能·pytorch·python·cann
Eloudy2 小时前
用 Python 直写 CUDA Kernel的技术,CuTile、TileLang、Triton 与 PyTorch 的深度融合实践
人工智能·pytorch
Rorsion4 小时前
PyTorch实现线性回归
人工智能·pytorch·线性回归
骇城迷影5 小时前
Makemore 核心面试题大汇总
人工智能·pytorch·python·深度学习·线性回归
mailangduoduo6 小时前
零基础教学连接远程服务器部署项目——VScode版本
服务器·pytorch·vscode·深度学习·ssh·gpu算力
多恩Stone7 小时前
【3D AICG 系列-6】OmniPart 训练流程梳理
人工智能·pytorch·算法·3d·aigc
前端摸鱼匠18 小时前
YOLOv8 环境配置全攻略:Python、PyTorch 与 CUDA 的和谐共生
人工智能·pytorch·python·yolo·目标检测
纤纡.1 天前
PyTorch 入门精讲:从框架选择到 MNIST 手写数字识别实战
人工智能·pytorch·python
子榆.1 天前
CANN 与主流 AI 框架集成:从 PyTorch/TensorFlow 到高效推理的无缝迁移指南
人工智能·pytorch·tensorflow