深度学习模型移植-替换torch.einsum算子

最新推荐文章于 2024-06-16 21:46:35 发布

好好学图算

最新推荐文章于 2024-06-16 21:46:35 发布

阅读量849

点赞数 2

文章标签：深度学习人工智能

本文链接：https://blog.csdn.net/u011231598/article/details/131933963

版权

文章介绍了当ONNX不支持torch.einsum算子时，如何使用基本的torch.matmul进行矩阵运算作为替代方法。通过重塑张量并应用matmul函数，可以在不支持einsum的嵌入式平台上实现相同的功能。这种方法适用于需要在限制硬件资源的平台上运行的深度学习模型。

摘要由CSDN通过智能技术生成

onnx不支持torch.einsum算子，很多嵌入式端平台就更不支持了，下面给出用基本的矩阵计算torch.matmul替代orch.einsum算子的代码。

B = 2
D = 3
H = 4
W1 = 5
W2 = 6

fmap1 = torch.randn(B, D, H, W1)
fmap2 = torch.randn(B, D, H, W2)
corr_einsum = torch.einsum('aijk,aijh->ajkh', fmap1, fmap2)
print(corr_einsum.shape) # torch.Size([2, 4, 5, 6])

fmap1 = fmap1.reshape(B*D*H,W1,1)
fmap2 = fmap2.reshape(B*D*H,1,W2)

corr = torch.matmul(fmap1, fmap2)
corr = corr.reshape(B,D,H,W1,W2)
corr = torch.sum(corr, dim=1)

print(corr.shape)
print(corr_einsum.equal(corr))

torch.matmul属于基本矩阵操作，一般嵌入式平台都会支持的，如果连基本的矩阵操作不支持，那就建议跟老板提出换芯片平台吧哈哈~