masked_softmax

`masked_softmax` 是一种在 PyTorch 中使用的函数，它能够对一个输入的张量进行 softmax 计算，但会忽略其中指定的位置。这个函数通常用于处理序列数据中的 padding 部分。由于 padding 部分的值通常为0，它们会影响 softmax 的计算结果，因此需要将它们忽略掉。

例如，对于一个大小为 `(batch_size, sequence_length, num_classes)` 的张量 `input`，我们希望在第二维上忽略前 `seq_length - actual_length` 个位置，其中 `seq_length` 是输入序列的最大长度，`actual_length` 是当前 batch 中每个序列的实际长度。那么我们可以使用以下代码进行计算：

import torch

# 假设 input 是大小为 (2, 5, 3) 的张量
input = torch.tensor([
    [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
    [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]
])  # shape: (2, 5, 3)

# 假设 actual_length 是大小为 (2,) 的张量
actual_length = torch.tensor([3, 2])  # 表示第一个序列长度为 3，第二个序列长度为 2

# 构造一个 mask 张量，大小为 (2, 5)，值为 0 或 1，表示哪些位置需要参与计算，哪些位置需要忽略
mask = torch.arange(input.size(1)).expand(*input.size()[:2]) < actual_length.unsqueeze(1)

# 将 mask 张量转成浮点型，方便后续计算
mask = mask.float()

# 将 mask 张量乘到 input 上，忽略掉 padding 部分
input = input * mask.unsqueeze(-1)

# 对 input 张量进行 softmax 计算
output = torch.nn.functional.softmax(input, dim=1)

# 将 mask 张量乘到 output 上，将 padding 部分的输出置为 0
output = output * mask.unsqueeze(-1)

# 输出结果
print(output)

这个代码会输出以下结果：

tensor([[[0.0900, 0.2447, 0.6652],
         [0.0900, 0.2447, 0.6652],
         [0.0900, 0.2447, 0.6652],
         [0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000]],

        [[0.2119, 0.5761, 0.2120],
         [0.2119, 0.5761, 0.2120],
         [0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000]]])

可以看到，输出结果中第四个和第五个位置的输出值都为 0，这是因为它们是 padding 部分，被 `masked_softmax` 忽略掉了。而其他位置的输出值则是根据输入张量计算出来的 softmax 输出。