首页 > 解决方案 > 如何从张量中获取列?

问题描述

假设我有一个由 1 和 0 组成的张量,如下所示。如何获取特定列的索引以替换为新值?如果我想用 [3.,4.,5.,6.] 替换第 1 列的值,我该如何完成呢?

a = torch.tensor([[[1., 0., 0., 0.]],
        [[0., 1., 0., 0.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]]])

标签: pytorchtensor

解决方案


考虑到这是一个 3D 张量,称它们为“列”有点棘手。

这将满足您的需要,将 'column' 1 设置为您提供的值。

a = torch.tensor([[[1., 0., 0., 0.]],
        [[0., 1., 0., 0.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]]])

# Change values in 'column' 1 (zero-indexed):
# The 0 is there because of the size-1 second dimension.
a[1, 0, :] = torch.tensor([3., 4., 5., 6.])

print(a)
# tensor([[[1., 0., 0., 0.]],
#         [[3., 4., 5., 6.]],
#         [[1., 0., 0., 0.]],
#         [[0., 0., 0., 1.]],
#         [[1., 0., 0., 0.]],
#         [[0., 0., 0., 1.]],
#         [[1., 0., 0., 0.]]])

推荐阅读