我有 2 个张量,A 和 B:
A = torch.randn([32,128,64,12],dtype=torch.float64)
B = torch.randn([64,12,64,12],dtype=torch.float64)
C = torch.tensordot(A,B,([2,3],[0,1]))
D = C.permute(0,2,1,3) # shape:[32,64,128,12]
张量 D 来自操作“tensordot -> permute”。如何实现新的操作 f() 以在 f() 之后进行张量操作,如下所示:
A_2 = f(A)
B_2 = f(B)
D = torch.tensordot(A_2,B_2)
最佳答案
您是否考虑过使用torch.einsum
哪个非常灵活?
D = torch.einsum('ijab,abkl->ikjl', A, B)
tensordot
的问题在于,它在 B
的维度之前输出 A
的所有维度以及您要查找的内容(排列时)是“交错”A
和 B
的维度。
关于pytorch - 排列后如何进行张量运算,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/65716540/