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 张量缩约”一个例子。

来源：51CTO

作者：lart

链接：https://blog.csdn.net/P_LarT/article/details/100156096