einsum算子不支持问题处理

原创已于 2024-07-15 19:39:28 修改 · 469 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch

于 2024-07-15 19:38:51 首次发布

当torch.einsum算子在onnx转om找不到时，如何进行代码替换。
torch.einsum里面为
bhwc, hkc->bhwk
计算方式为bhwc乘以hkc得到bhwk

b h w c
  h k c
b h w k
   ||
   \/
b h w c
  h c k(交换c和k的位置，bhwc*ck=bhwk)
b h w k

举例：

M = torch.randn(100, 30, 40, 60)
N = torch.randn(30, 40, 60)
torch_transform_result = torch.einsum("bhwc,hkc->bhwk", M, N)

## 可以替换为torch.matmul
N_t = torch.transpose(N, 1, 2)
result = torch.matmul(M, N_t) ## 此结果和torch_transform_result是一样的。