【Pytorch学习笔记】4.细讲Pytorch的gather函数是什么——从Softmax回归中交叉熵损失函数定义的角度讲述

本文接着慢慢磨pytorch基础。本来是想记录一下心得，结果码着码着又讲成了story。

gather函数：原始数据矩阵 根据索引矩阵 取到对应值矩阵

我们在学习Softmax回归从零实现的时候，需要定义一个交叉熵损失函数。我们会使用torch.gather函数的方法取原始数据矩阵中对应位置的值，接着再取log等处理。

可以先看看的例子

可以先囫囵吞枣地看一下gather函数的例子：

# 变量y_hat是2个样本在3个类别的预测概率，变量y是这2个样本的标签类别。
import torch
y_hat = torch.tensor([[0.1, 0.3, 0.6], [0.3, 0.2, 0.5]])
# y是一个索引矩阵，包含了每个样本的正确类别的位置。
# 比如第1个样本是第1类，第2个样本是第3类。0即第1个，2即第2个。
y = torch.LongTensor([0, 2])
y_hat.gather(1, y.view(-1, 1))

# 输出：
tensor([[0.1000],
        [0.5000]])

可以看到y是一个索引矩阵，包含了每个样本的正确类别的位置。比如第1个样本是第1类，第2个样本是第3类（0是第1个，2是第3个）。
原始数据矩阵y_hat使用gather传入y后，成功取到了两个样本的预测概率分布 各自对应正确类别下的 概率。
所以这里已经可以理解：gather函数是 原始数据矩阵 根据 索引矩阵 取到 对应值矩阵 的一个过程。

这样我们就可以定义交叉熵损失函数：

def cross_entropy(y_hat, y):
    return - torch.log(y_hat.gather(1, y.view(-1, 1)))

官方文档解释

官方文档中对gather函数的描述非常简洁，即按设定的维度方向，按该维度上的索引值取值。

给的例子也很直白。但是直观上看这是一个逻辑上非常绕的函数，按维度方向根据索引取值也是很难想象的。

我的思考是为什么我们要用这么别扭的方法取值呢？定义交叉熵损失函数为啥要用gather呢？

从交叉熵损失函数定义的角度理解gather函数的使用

1. 在Softmax回归中，我们知道要使用交叉熵损失函数来计算两个概率分布之间的差异。
   还是一开始的例子，我们有 2个样本3类别的预测概率分布y_hat 和实际的概率分布y，如图：
   

2. 根据交叉熵损失函数：
   
   对于softmax回归来说，真实的概率分布只有在正确的类别值为1，所以损失函数最后可以化简为：
   $-\log {\hat{y}}_p^{(i)}$。即在预测概率分布中取到第p个值，而p就是索引值，对应正确的类别。

3. 对于一个batch的数据来说，每一行是一个样本的各类别概率分布，每一行都会有一个正确分类，位于第$p^{(i)}$列。这个第$p^{(i)}$列也是真实分布y每一行中1的位置。所以我们只要知道每个样本正确的类别所在列的标号就行了。
   即对于实际概率分布y矩阵：
   我们只需要把y矩阵转变为索引矩阵。
   故对 y_hat 的每个样本取索引值就得到了对应类别的预测概率。之后就可以进行取log等操作来定义损失函数了。

4. 这样，根据一个索引矩阵对一个原始数据矩阵的行（或列）取索引对应的值，就是gather函数的具体做的事了。而交叉熵损失函数的定义正好符合了这样的数据特征，所以就能正好使用gather函数了。
   （这里的索引矩阵y在实际例子中就是一个batch数据X的对应label y。）

这样做有什么好处呢？个人觉得：
6. 将真实概率矩阵（抑或叫One-Hot编码矩阵）缩减为索引矩阵可以大大减少内存开销。
7. 索引值矩阵的维数没有变化，依然保留了矩阵并行运算效率高的优点。

torch.gather()

讲了半天交叉熵损失函数，最后讲讲gather函数的使用。
对于初学者来说，最重要的参数就是dim。在开头定义交叉熵损失函数的时候y_hat.gather(1, y.view(-1, 1))，参数dim=1，以行为方向进行索引。
直接用二维Tensor举例：

import torch
src = torch.arange(1, 16).reshape(5, 3)
"""
src:
tensor([[ 1,  2,  3],
        [ 4,  5,  6],
        [ 7,  8,  9],
        [10, 11, 12],
        [13, 14, 15]])
"""

# 定义两个索引矩阵
index1 = torch.tensor([[0, 1, 2], [2, 3, 4], [0, 2, 4]])
index2 = torch.tensor([[1, 2, 0, 2, 1], [1, 2, 1, 0, 0]]).t()
"""
index1:
tensor([[0, 1, 2],
        [2, 3, 4],
        [0, 2, 4]])
index2:
tensor([[1, 1],
        [2, 2],
        [0, 1],
        [2, 0],
        [1, 0]])
"""

# axis=0时，
output1 = src.gather(dim=0, index=index1)
print(output1)
"""
输出:
tensor([[ 1,  5,  9],
        [ 7, 11, 15],
        [ 1,  8, 15]])
"""

# axis=1
output2 = src.gather(dim=1, index=index2)
print(output2)
"""
输出:
tensor([[ 2,  2],
        [ 6,  6],
        [ 7,  8],
        [12, 10],
        [14, 13]])
"""

当一个数据矩阵用一个索引矩阵在某个维度方向上取值，那个方向上的所有值 按 索引矩阵中那个方向的索引值 来取值。

我们直接看下图就更加清晰明了（原图中Dim以1、2来说明，在pytorch中即为0、1）：

指定了一个方向以后，index以指定方向一刀切下去，对这个方向的数字取值。
index和src的维度必须要一致，且除去切的那一维，其余维度shape也一致（扩展到三维也是这样）。