pytorch的einsum

# -*- coding: utf-8 -*-
import torch

# 转置
a = torch.arange(6).reshape(2, 3)
print('a:{}'.format(a))
b = torch.einsum('ij->ji', [a])
print(b.size())
print(b)

# 求和
b = torch.einsum('ij->', [a])
print(b.size(), b)
# 行求和
b = torch.einsum('ij->i', [a])
print(b.size(), b)
# 列求和
b = torch.einsum('ij->j', [a])
print(b.size(), b)

# 矩阵求积
d = torch.arange(12).reshape(3, 4)
print(f'd:{d}')
e = torch.einsum('ik,kh->ih', [a, d])
print(e.size(), e)

# batch
a = torch.randn(3, 2, 5)
b = torch.randn(3, 5, 3)
c = torch.einsum('ijk,ikl->ijl', [a, b])
print(c.size())
# 张量reduce
a = torch.arange(120).reshape(2, 3, 4, 5)
b = torch.arange(10).reshape(2, 5)
c = torch.einsum('bijh,rh->birj', [a, b])
print(a)
print(b)
print(c)
print(c.shape)

参考：
https://zhuanlan.zhihu.com/p/44954540

最后

以上就是清爽小土豆为你收集整理的pytorch的einsum的全部内容，希望文章能够帮你解决pytorch的einsum所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。