PyTorch之爱因斯坦求和约定

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

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

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

先写一个一般的思路。首先要注意PyTorch中的维度顺序为N, C, H, W学习PyTorch的关键是要记住这个顺序。对于原文来说，下表s、t、i、j分别表示长、宽、通道、通道，所以对于这里提到的张量$F_{s,t,i}^1$和$F_{s,t,j}^2$各自实际上是对应于形状为[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 张量缩约”一个例子。

最后

以上就是懦弱钢笔为你收集整理的PyTorch之爱因斯坦求和约定PyTorch之爱因斯坦求和约定的全部内容，希望文章能够帮你解决PyTorch之爱因斯坦求和约定PyTorch之爱因斯坦求和约定所遇到的程序开发问题。

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