PyTorch中如何高效切片张量并按指定规则修改其值？

问题背景 #

给定三个PyTorch张量：

import torch
a = torch.tensor([[5, 2, 3, 24],
                  [8, 66, 7, 89],
                  [9, 10, 1, 12]])

b = torch.tensor([[10, 22, 13, 1],
                  [35, 6, 17, 3],
                  [11, 13, 5, 8]])

c = torch.tensor([[0, 0, 0, 0],
                  [0, 0, 0, 0],
                  [0, 0, 0, 0]])

需要完成两步操作：

1. 先修改c的前3列（除最后一列外），计算公式为：c[:,:-1] = a[:,:-1] - a[:,1:] - b[:,1:]，计算后c的中间状态为：

tensor([[-19, -14, -22, 0],
        [-64, 42, -85, 0],
        [-14, 4, -19, 0]])

2. 最后修改c的最后一列：c[:, -1] = b[:, -1] - 1，得到最终目标张量。

询问：在PyTorch中实现该张量切片与修改的最快方法是什么？

解答 #

在PyTorch里，直接使用原地切片赋值就是效率最高的方式，这种操作不需要额外创建新张量，直接在原张量c的内存空间上修改数据，避免了不必要的内存分配和数据拷贝开销。

完整实现代码如下：

import torch

# 定义初始张量
a = torch.tensor([[5, 2, 3, 24],
                  [8, 66, 7, 89],
                  [9, 10, 1, 12]])

b = torch.tensor([[10, 22, 13, 1],
                  [35, 6, 17, 3],
                  [11, 13, 5, 8]])

c = torch.tensor([[0, 0, 0, 0],
                  [0, 0, 0, 0],
                  [0, 0, 0, 0]])

# 第一步：修改前3列
c[:, :-1] = a[:, :-1] - a[:, 1:] - b[:, 1:]

# 第二步：修改最后一列
c[:, -1] = b[:, -1] - 1

# 查看最终结果
print(c)

执行后输出的最终张量为：

tensor([[-19, -14, -22,  0],
        [-64,  42, -85,  2],
        [-14,   4, -19,  7]])

为什么这是最快的方式？ #

PyTorch的切片赋值属于原地操作（in-place operation），对c的切片直接赋值时，会直接覆盖原张量对应位置的内存数据，不需要为中间结果分配新内存块，也不需要把数据从新张量拷贝回c，在处理大规模张量时优势尤为明显。如果先计算结果再整体拼接赋值，反而会额外创建多个临时张量，增加内存占用和计算耗时。

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