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 格式。
- 更新张量特定位置的值:可以根据索引在张量的某些特定位置填入新的值。
它灵活且高效,适合用于需要对张量的特定索引进行操作的场景。