onnx中的where op

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])

torch.onnx.export(model, im, "new.onnx", do_constant_folding=True, verbose=False, opset_version=11, 
                  input_names=["x"],output_names=["out"])

onnx_model = onnx.load("new.onnx")

model_simp, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model_simp, "new_sim.onnx")

print('finished exporting simplified onnx')

简化前：

简化后：

参考： 

Performance degradation on GatherElements · Issue #4119 · onnx/onnx · GitHub

PyTorch学习笔记——repeat()和expand()区别 - 知乎