onnx中的where op

曙光_deeplove

已于 2022-04-21 11:31:58 修改

阅读量2.5k

点赞数

CC 4.0 BY-SA版权

于 2022-04-21 09:22:00 首次发布

本文链接：https://blog.youkuaiyun.com/yangjinyi1314/article/details/124313248

深度学习算法部署同时被 3 个专栏收录

28 篇文章 ¥49.90 ¥99.00

订阅专栏

超级会员免费看

ONNX

12 篇文章 ¥69.90 ¥99.00

订阅专栏

超级会员免费看

pytorch

14 篇文章

订阅专栏

博客探讨了在ONNX中where操作如何由PyTorch的特定操作如torch.where和expand产生。文章指出，torch.where在特定调用方式下会产生where节点，而expand通常会导致where节点的出现。作者还提到了在不同TRT版本中where操作的转换问题，并提供了相关源码和链接以便进一步研究。

onnx中的where op是pytorch中的什么操作产生的呢？试了torch.where 以及 a[mask]等操作都没有产生。也没搜到。想记录学习下。

答案：看下expand （我自己给的答案），我还没有验证

球哥给的答案：

expand可以有where。单独的torch.where也会有where节点并且只有在 torch.where(a<0, x, y)这种调用下会有where

看这个源码就能看出来，否则就会转为其他节点！

trt7.2.1之前where转的时候好像报错，但是高版本没问题。

import torch
import onnx
from onnxsim import simplify

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, x):
        x = x.expand(2, 3)
        return x

model = Model()

im = torch.tensor([1, 2, 3])

了解本专栏