torch.einsum函数可以实现多种张量运算,如矩阵乘法、向量点积、外积、迹、对角化等。
基本语法:
torch.einsum(equation, *operands)
equation
是一个描述运算的字符串,格式为'inputs->outputs'
,其中inputs
描述输入张量的维度,outputs
描述输出张量的维度。*operands
是一个或多个输入张量。
基本用法
- 矩阵乘法
假设有两个矩阵 (A) 和 (B),维度分别为 (m x n) 和 (n x p)。矩阵乘法的 einsum
表示为:
C = torch.einsum('ij,jk->ik', A, B)
ij
代表矩阵 (A) 的两个维度,jk
代表矩阵 (B) 的两个维度,ik
代表输出矩阵 (C) 的维度。
- 向量点积
对于两个向量 (a) 和 (b),其点积可以表示为:
dot_product = torch.einsum('i,i->', a, b)
- 外积
对于两个向量 (a) 和 (b),其外积(即生成一个矩阵,其中每个元素是 )为:
outer_product = torch.einsum('i,j->ij', a, b)
- 迹
对于一个矩阵 (A),其迹(即对角线元素之和)为:
trace = torch.einsum('ii->', A)
- 对角化
将一个向量 (v) 转换为一个对角矩阵,其对角线元素为 (v) 的元素:
diag_matrix = torch.einsum('i,j->ij', v, torch.arange(len(v)) == torch.arange(len(v))[:, None])
diag_matrix = torch.einsum('i,j->ij', v, (torch.eye(len(v)) == 1).float())
#这里使用 (torch.eye(len(v)) == 1).float() 是为了生成一个对角线上全为 1,其余位置为 0 的浮点矩阵。
- 批量矩阵乘法
假设你有两个三维张量 (A) 和 (B),维度分别为 (b x m x n) 和 (b x n x p),其中 (b) 是批量大小。批量矩阵乘法的 einsum
表示为:
C = torch.einsum('bij,bjk->bik', A, B)
- 多矩阵乘法
假设有三个矩阵 (A), (B), 和 (C),维度分别为 (l x m), (m x n), 和 (n x o),计算 (ABC):
D = torch.einsum('ij,jk,kl->il', A, B, C)