【踩坑】pytorch中的索引与copy_结合不会复制数据及其解决方案

转载请注明出处：小锋学长生活大爆炸[xfxuezhagn.cn]

如果本文帮助到了你，欢迎[点赞、收藏、关注]哦~

目录

背景知识

实验验证

结论分析

错误案例

处理方法

注意事项

附加说明

基本索引返回视图

高级索引返回副本

赋值操作都是原地操作

---

以下内容若有错误，欢迎指出！

背景知识

        我们现在知道以下几个知识点：

1. copy_函数：pytorch中变量的copy_函数，可以将变量inplace地复制到另一个变量中。比如buffer.copy_(a)，将a中的数据直接复制到了buffer中。
2. 视图(View)： 视图是指不复制数据的情况下，返回一个与原始张量共享内存的张量。例如，切片操作通常会返回一个视图。
3. 副本(Copy)： 副本是指返回一个新的张量，包含了原始张量的数据，但不共享内存。布尔掩码索引返回的就是这样的副本。
4. PyTorch和Numpy中的情况：在通过索引访问张量的内容时，PyTorch 遵循 Numpy 的行为，即基本索引返回视图，而高级索引返回副本。通过基本索引或高级索引进行的赋值都是原地操作。
5. 切片()：切片由中括号和冒号组成，如[:10]、[2:10]、[2:]。
6. 基本索引：使用整数或切片来访问数组的元素。
7. 高级索引：指的是使用整数数组、布尔数组或者其他序列来访问数组的元素。相比于基本索引，高级索引可以访问到数组中的任意元素，并且可以用来对数组进行复杂的操作和修改。

来自：PyTorch入门笔记-索引和切片

        好了，现在有一个问题，如果结合索引与copy_操作，那是否会复制成功？

实验验证

        答案是，不会成功。我们可以用代码测试一下：

import torchprint('>> 使用=号直接复制 <<')
buff = torch.arange(5)
mask = [True, False, False, True, False]
print('输出原始变量：', buff)
print('输出索引掩码：', mask)
print('输出索引变量：', buff[mask])
buff[mask][0] = 10
print('索引变量修改：', buff)
buff[mask] = torch.tensor([8, 9])
print('索引变量赋值：', buff)print('*' * 50)buff = torch.arange(5)
print('输出原始变量：', buff)
print('输出切片索引：', '1:3')
buff_indices = buff[1:3]
print('输出切片变量：', buff[buff_indices])
buff[1:3][0] = 10
print('切片变量修改：', buff)
buff[1:3] = torch.tensor([8, 9])
print('切片变量赋值：', buff)print('=' * 50)print('>> 使用copy_原地复制 <<')buff = torch.arange(5)
mask = [True, False, False, True, False]
print('输出原始变量：', buff)
print('输出索引掩码：', mask)
print('输出索引变量：', buff[mask])
buff[mask].copy_(torch.tensor([8, 9]))
print('索引变量copy：', buff)print('*' * 50)buff = torch.arange(5)
print('输出原始变量：', buff)
print('输出切片索引：', '1:3')
print('输出切片变量：', buff[1:3])
buff[1:3].copy_(torch.tensor([8, 9]))
print('切片变量copy：', buff)

输出结果(改变的地方加粗了)：

>> 使用=号直接复制 <<
输出原始变量： tensor([0, 1, 2, 3, 4])
输出索引掩码： [True, False, False, True, False]
输出索引变量： tensor([0, 3])
索引变量修改： tensor([0, 1, 2, 3, 4])
索引变量赋值： tensor([8, 1, 2, 9, 4])
**************************************************
输出原始变量： tensor([0, 1, 2, 3, 4])
输出切片索引： 1:3
输出切片变量： tensor([1, 2])
切片变量修改： tensor([ 0, 10,  2,  3,  4])
切片变量赋值： tensor([0, 8, 9, 3, 4])
==================================================
>> 使用copy_原地复制 <<
输出原始变量： tensor([0, 1, 2, 3, 4])
输出索引掩码： [True, False, False, True, False]
输出索引变量： tensor([0, 3])
索引变量赋值： tensor([0, 1, 2, 3, 4])
**************************************************
输出原始变量： tensor([0, 1, 2, 3, 4])
输出切片索引： 1:3
输出切片变量： tensor([1, 2])
切片变量赋值： tensor([0, 8, 9, 3, 4])

结论分析

        在PyTorch中，当你使用布尔掩码或索引来访问张量时，通常会创建一个新的张量，而不是对原始张量进行原地修改。在PyTorch中，切片操作通常会返回一个视图，而不是数据的副本。这意味着切片操作返回的张量和原始张量共享相同的内存。因此，对切片后的张量进行的任何修改都会影响到原始张量。与此相对，布尔掩码索引返回的是数据的副本，因此修改索引得到的张量不会影响原始张量。

        因此可见，由于索引返回的是新张量，而copy_是原地复制，因此对于原来的变量来说并没有影响，所以不会复制成功。

        而=号这个赋值操作，不管是基本索引还是高级索引，由于底层都是对张量的原地操作，因此确实可以赋值成功。

错误案例

        根据以上内容就知道，有时候我们如果这样用，那就是错的：

buff = torch.arange(5)
mask = [True, False, False, True, False]
buff[mask].copy_(torch.tensor([8, 9]))

处理方法

        如果确实想结合索引和copy_一起用怎么办？那么可以试试masked_scatter_。

torch.Tensor.masked_scatter_ — PyTorch 2.4 documentation

注意事项

1. 对于 [1,2,3] 这种，仍然是位置索引，并不是切片，所以效果等同于上面的布尔索引。通常，PyTorch中的张量索引使用逗号分隔的整数索引来指定每个维度上的具体位置。如果你想对一个一维张量进行切片，应该使用冒号:来指定范围。

附加说明

        对于背景知识里的第4点，我们也来通过代码验证一下。

基本索引返回视图

基本索引包括标量索引、切片操作和整数索引。PyTorch通常会返回原始张量的视图，这意味着它们共享相同的底层数据。因此，对视图的修改会影响原始张量。例如：

import torcha = torch.tensor([1, 2, 3, 4])
b = a[:2]  # 基本索引，b 是 a 的视图
b[0] = 10  # 修改视图会影响原始张量
print(a)   # 输出: tensor([10,  2,  3,  4])

高级索引返回副本

高级索引包括使用布尔数组、整数数组或多维索引。PyTorch和NumPy一样，高级索引会返回一个新的张量，即副本，不与原始数据共享内存。因此，对副本的修改不会影响原始张量。例如：

import torcha = torch.tensor([1, 2, 3, 4])
indices = torch.tensor([0, 2])
b = a[indices]  # 高级索引，b 是 a 的副本
b[0] = 10       # 修改副本不会影响原始张量
print(a)        # 输出: tensor([1, 2, 3, 4])
print(b)        # 输出: tensor([10,  3])

赋值操作都是原地操作

无论是通过基本索引还是高级索引，赋值操作都是原地操作，这意味着它们会直接修改原始张量的内容。例如：

1. 基本索引赋值：

   a = torch.tensor([1, 2, 3, 4])
   a[:2] = torch.tensor([10, 20])  # 原地修改 a
   print(a)  # 输出: tensor([10, 20,  3,  4])
   

2. 高级索引赋值：

   a = torch.tensor([1, 2, 3, 4])
   indices = torch.tensor([0, 2])
   a[indices] = torch.tensor([10, 20])  # 原地修改 a
   print(a)  # 输出: tensor([10,  2, 20,  4])