pytorch scatter_ 函数介绍

scatter_ 是 PyTorch 中的一个原地操作函数，用于在给定的索引处将某些值填充到张量的指定维度中。它的常见用途之一是将类别标签转换为 one-hot 编码，不过它也适用于其他场景，如在特定索引处更新张量的值。

scatter_ 函数的签名如下：

scatter_(dim, index, src)

• dim：指定操作的维度。即在这个维度上更新值。
• index：包含索引的张量，指定要更新值的位置。
• src：要填入的值。可以是一个标量（单个值），也可以是一个张量。

使用示例

1. 使用 scatter_ 实现 one-hot 编码

我们可以通过 scatter_ 将类别标签转换为 one-hot 编码。

代码示例：

import torch

# logits 模拟网络输出 [batch_size, num_classes]
logits = torch.tensor([[2.0, 1.0, 0.1], [1.5, 2.5, 0.5]])  # 形状 [2, 3]

# target 是真实的类别标签 [batch_size]
target = torch.tensor([0, 2])  # 形状 [2]

# 创建一个与 logits 相同大小的全零张量
target_onehot = torch.zeros_like(logits)  # 形状 [2, 3]

# 使用 scatter_ 函数在第 1 维（类别维度）根据 target 的索引设置为 1
target_onehot.scatter_(1, target.view(-1, 1), 1)

print("one-hot 编码：")
print(target_onehot)

注：target.view(-1, 1) 中的 1 指的是将 target 张量的形状重新调整为 两维 ，并且使得第二个维度的大小固定为 1。

输出：

tensor([[1., 0., 0.],  # 第一个样本，类别为 0
        [0., 0., 1.]])  # 第二个样本，类别为 2

2. 使用 scatter_ 更新指定位置的值

你还可以用 scatter_ 在张量的指定位置填充任意值。这里是一个简单的例子，将特定索引的位置设置为自定义的数值：

代码示例：

import torch

# 创建一个 3x3 的全零张量
tensor = torch.zeros(3, 3)

# 定义索引
index = torch.tensor([[0, 2, 1]])  # 每行对应位置的索引

# 要填入的值
src = torch.tensor([[5, 9, 7]])

# 使用 scatter_ 在第 1 维（列）填充 src 的值到指定索引位置
tensor.scatter_(1, index, src)

print("填充值后的张量：")
print(tensor)

输出：

tensor([[5., 7., 9.],  # 在索引 0 处填 5，索引 1 处填 7，索引 2 处填 9
        [5., 7., 9.],
        [5., 7., 9.]])

3. scatter_ 与广播

scatter_ 支持广播机制。你可以使用一个标量值来替换指定的索引位置，也可以使用一个与 index 兼容的张量来填充不同的值。

代码示例：

import torch

# 创建一个 4x3 的全零张量
tensor = torch.zeros(4, 3)

# 定义索引
index = torch.tensor([[0, 2, 1], [2, 1, 0], [1, 0, 2], [2, 1, 0]])  # 4x3 的索引

# 要填入的值
src = torch.tensor([5, 9, 7])  # 广播到每一行

# 使用 scatter_ 在第 1 维（列）填充 src 的值到指定索引位置
tensor.scatter_(1, index, src)

print("填充值后的张量：")
print(tensor)

输出：

tensor([[5., 7., 9.],
        [7., 9., 5.],
        [9., 5., 7.],
        [7., 9., 5.]])

总结

scatter_ 函数可以根据指定的索引，在目标张量的某个维度上填充源张量或标量的值。它的常见应用包括：

• one-hot 编码：将类别标签转换为 one-hot 格式。
• 更新张量特定位置的值：可以根据索引在张量的某些特定位置填入新的值。

它灵活且高效，适合用于需要对张量的特定索引进行操作的场景。