PyTorch Tensor.scatter_方法dim参数作用及独热编码使用问题

scatter_方法dim参数的核心工作逻辑 #

Tensor.scatter_(dim, index, src)是PyTorch的原地赋值操作，作用是跟随索引张量index的指引，把src中的值填充到当前张量的对应位置，其中dim参数决定了索引替换的维度方向，以最常用的二维张量场景为例，规则非常清晰：

• 当传入dim=0时，索引替换发生在第0维度（行方向）：对于index中任意坐标为(r, c)的元素，其存储的索引值v会作为目标位置的行号，最终赋值位置为(v, c)，也就是固定列号不变，沿着行方向上下移动找填充位置。
• 当传入dim=1时，索引替换发生在第1维度（列方向）：对于index中任意坐标为(r, c)的元素，其存储的索引值v会作为目标位置的列号，最终赋值位置为(r, v)，也就是固定行号不变，沿着列方向左右移动找填充位置。

独热编码场景选择dim=1的原因 #

问题中给出的实现代码如下：

>>> target = torch.tensor([3, 5, 0, 2, 7, 5])
>>> target
tensor([3, 5, 0, 2, 7, 5])

>>> onehot = torch.zeros(target.shape[0], 8)
>>> onehot.scatter_(1, target.unsqueeze(1), 1.0)
tensor([[0., 0., 0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 1., 0., 0.],
        [1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 0., 1., 0., 0.]])

做分类任务的独热编码时，我们的预期输出是形状为(样本数, 类别总数)的张量，每一行对应一个样本，行内只有该样本类别对应的列位置为1，其余位置为0。
这个需求和dim=1的赋值逻辑完全匹配：

• 初始化的onehot形状为(6,8)，对应6个样本、共8个可能类别
• 传入的索引张量是target.unsqueeze(1)，形状为(6,1)，每个元素的坐标为(样本序号, 0)，存储的值是对应样本的类别标签
• 设置dim=1时，每个索引元素存储的类别标签会被作为列号，固定行号为样本序号，刚好把1填充到对应行的对应类别列，直接得到预期的独热结构。

如果设置dim=0，所有的1都会被填充到第0列的不同行位置，最终得到的张量只有第0列零散存在1，完全不符合独热编码每行一个有效位的要求。

内容的提问来源于stack exchange，提问作者NewHere