PyTorch之爱因斯坦求和约定

PyTorch之爱因斯坦求和约定

文章目录

• ​​PyTorch之爱因斯坦求和约定​​

• ​​关于`einsum`维度记忆的小技巧​​

网上关于这个函数：

​​torch.einsum​​的介绍已经很多了，这里列出我重点看过的一篇文章。

• ​​einsum满足你一切需要：深度学习中的爱因斯坦求和约定​​

这篇文章写的非常棒，很详细。

这里写个简单的例子，对于论文​​A Gift from Knowledge Distillation: Fast Optimization, Network Minimization and Transfer Learning​​中的下面的式子，可以很方便的借助该函数搞定。

先写一个一般的思路。首先要注意PyTorch中的维度顺序为​​N, C, H, W​​​学习PyTorch的关键是要记住这个顺序。对于原文来说，下表s、t、i、j分别表示长、宽、通道、通道，所以对于这里提到的张量$F^1_{s,t,i}$和$F^2_{s,t,j}$各自实际上是对应于形状为​​​[b, m, h, w]​​​和​​[b, n, h, w]​​​的。而这里的累加符号是对于$s$和$t$进行的计算，所以实际上可以转化为矩阵乘法。​​​[b, m, hxw] * [b, hxw, n] = [b, m, n]​​

先准备数据：

import torch
a = torch.rand(2, 3, 4, 5)
b = torch.rand_like(a)

一般的利用矩阵乘法的思路：

c = torch.bmm(a.view(2, 3, -1), b.view(2, 3, -1).transpose(1, 2)) / (4 * 5)

而当使用​​torch.einsum​​的时候只需要一行：

d = torch.einsum("bist,bjst->bij", [a, b]) / (4 * 5)

验证结果：

d == c

# output：
tensor([[[1, 1, 1],
         [1, 1, 1],
         [1, 1, 1]],
        [[1, 1, 1],
         [1, 1, 1],
         [1, 1, 1]]], dtype=torch.uint8)

关于​​einsum​​维度记忆的小技巧

以前面的​​torch.einsum("bist,bjst->bij", [a, b])​​​为例。这里​​einsum​​的第一个参数表示的维度变换关系，也就是各个维度的自己的索引下标。一般只要满足对应关系即可。即各个维度使用不同的下标，如果一样，那就会一起进行加和计算（可以认为是对于索引遍历的过程汇总二者是同步的）。对于逗号的分隔表示对应于后面​​[]​​​（也可以不用​​[]​​​包裹，因为第二部分参数使用的是一个可变长参数接收的）中的不同张量，也就是这里的​​bist​​​对应于​​a​​​，而​​bjst​​​对应于​​b​​​。下标按照张量的对应维度调整好之后，就可以开始计算了。由于这里计算的是针对​​s&t​​​的累加，最后​​s&t​​​都消除了，仅剩下来​​d​​​的​​b&i&j​​​三个索引。所以也就顺其自然的写出了这样的变换关系：​​bist,bjst->bij​​。相当的方便！

这里实际上就是前面参考文章中的“2.11 张量缩约”一个例子。