PyTorch中的拷贝与就地操作详解_Python

前言

PyTroch中我们经常使用到Numpy进行数据的处理，然后再转为Tensor，但是关系到数据的更改时我们要注意方法是否是共享地址，这关系到整个网络的更新。本篇就In-palce操作，拷贝操作中的注意点进行总结。

In-place操作

pytorch中原地操作的后缀为_，如.add_()或.scatter_()，就地操作是直接更改给定Tensor的内容而不进行复制的操作，即不会为变量分配新的内存。Python操作类似+=或*=也是就地操作。（我加了我自己~）

为什么in-place操作可以在处理高维数据时可以帮助减少内存使用呢，下面使用一个例子进行说明，定义以下简单函数来测量PyTorch的异位ReLU（out-of-place）和就地ReLU(in-place)分配的内存：

				?

									import torch # import main library

									import torch.nn as nn # import modules like nn.ReLU()

									import torch.nn.functional as F # import torch functions like F.relu() and F.relu_()

									def get_memory_allocated(device, inplace = False):

									 '''

									 Function measures allocated memory before and after the ReLU function call.

									 INPUT:

									 - device: gpu device to run the operation

									 - inplace: True - to run ReLU in-place, False - for normal ReLU call

									 '''

									 # Create a large tensor

									 t = torch.randn(10000, 10000, device=device)

									 # Measure allocated memory

									 torch.cuda.synchronize()

									 start_max_memory = torch.cuda.max_memory_allocated() / 1024**2

									 start_memory = torch.cuda.memory_allocated() / 1024**2

									 # Call in-place or normal ReLU

									 if inplace:

									 F.relu_(t)

									 else:

									 output = F.relu(t)

									 # Measure allocated memory after the call

									 torch.cuda.synchronize()

									 end_max_memory = torch.cuda.max_memory_allocated() / 1024**2

									 end_memory = torch.cuda.memory_allocated() / 1024**2

									 # Return amount of memory allocated for ReLU call

									 return end_memory - start_memory, end_max_memory - start_max_memory

									 # setup the device

									device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")

									#开始测试

									# Call the function to measure the allocated memory for the out-of-place ReLU

									memory_allocated, max_memory_allocated = get_memory_allocated(device, inplace = False)

									print('Allocated memory: {}'.format(memory_allocated))

									print('Allocated max memory: {}'.format(max_memory_allocated))

									'''

									Allocated memory: 382.0

									Allocated max memory: 382.0

									'''

									#Then call the in-place ReLU as follows:

									memory_allocated_inplace, max_memory_allocated_inplace = get_memory_allocated(device, inplace = True)

									print('Allocated memory: {}'.format(memory_allocated_inplace))

									print('Allocated max memory: {}'.format(max_memory_allocated_inplace))

									'''

									Allocated memory: 0.0

									Allocated max memory: 0.0

									'''

看起来，使用就地操作可以帮助我们节省一些GPU内存。但是，在使用就地操作时应该格外谨慎。

就地操作的主要缺点主要原因有2点，官方文档：

1.可能会覆盖计算梯度所需的值，这意味着破坏了模型的训练过程。

2.每个就地操作实际上都需要实现来重写计算图。异地操作Out-of-place分配新对象并保留对旧图的引用，而就地操作则需要更改表示此操作的函数的所有输入的创建者。

在Autograd中支持就地操作很困难，并且在大多数情况下不鼓励使用。Autograd积极的缓冲区释放和重用使其非常高效，就地操作实际上降低内存使用量的情况很少。除非在沉重的内存压力下运行，否则可能永远不需要使用它们。

总结：Autograd很香了，就地操作要慎用。

拷贝方法

浅拷贝方法：共享 data 的内存地址，数据会同步变化

* a.numpy() # Tensor—>Numpy array

* view() #改变tensor的形状，但共享数据内存，不要直接使用id进行判断

* y = x[:] # 索引

* torch.from_numpy() # Numpy array—>Tensor

* torch.detach() # 新的tensor会脱离计算图，不会牵扯梯度计算。

* model:forward()

还有很多选择函数也是数据共享内存，如index_select() masked_select() gather()。

以及后文提到的就地操作in-place。

深拷贝方法：

* torch.clone() # 新的tensor会保留在计算图中，参与梯度计算

下面进行验证，首先验证浅拷贝：

				?

									import torch as t

									import numpy as np

									a = np.ones(4)

									b = t.from_numpy(a) # Numpy->Tensor

									print(a)

									print(b)

									'''输出：

									[1. 1. 1. 1.]

									tensor([1., 1., 1., 1.], dtype=torch.float64)

									'''

									b.add_(1)# add_会修改b自身

									print(a)

									print(b)

									'''输出：

									[2. 2. 2. 2.]

									tensor([2., 2., 2., 2.], dtype=torch.float64)

									b进行add操作后, a,b同步发生了变化

									'''

Tensor和numpy对象共享内存（浅拷贝操作），所以他们之间的转换很快，且会同步变化。

造torch中y = x + y这样的运算是会新开内存的，然后将y指向新内存。为了进行验证，我们可以使用Python自带的id函数：如果两个实例的ID一致，那么它们所对应的内存地址相同；但需要注意是在torch中还有些特殊，数据共享时直接打印tensor的id仍然会出现不同。

				?

									x = torch.tensor([1, 2])

									y = torch.tensor([3, 4])

									id_0 = id(y)

									y = y + x

									print(id(y) == id_0) 

									# False

这时使用索引操作不会开辟新的内存，而想指定结果到原来的y的内存，我们可以使用索引来进行替换操作。比如把x + y的结果通过[:]写进y对应的内存中。

				?

									x = torch.tensor([1, 2])

									y = torch.tensor([3, 4])

									id_0 = id(y)

									y[:] = y + x

									print(id(y) == id_0) 

									# True

另外，以下两种方式也可以索引到相同的内存：

torch.add(x, y, out=y)
y += x, y.add_(x)

				?

									x = torch.tensor([1, 2])

									y = torch.tensor([3, 4])

									id_0 = id(y)

									torch.add(x, y, out=y) 

									# y += x, y.add_(x)

									print(id(y) == id_0) 

									# True

clone() 与 detach() 对比

Torch 为了提高速度，向量或是矩阵的赋值是指向同一内存的，这不同于 Matlab。如果需要保存旧的tensor即需要开辟新的存储地址而不是引用，可以用 clone() 进行深拷贝，

首先我们来打印出来clone()操作后的数据类型定义变化：

(1). 简单打印类型

				?

									import torch

									a = torch.tensor(1.0, requires_grad=True)

									b = a.clone()

									c = a.detach()

									a.data *= 3

									b += 1

									print(a) # tensor(3., requires_grad=True)

									print(b)

									print(c)

									'''

									输出结果：

									tensor(3., requires_grad=True)

									tensor(2., grad_fn=<AddBackward0>)

									tensor(3.)  # detach()后的值随着a的变化出现变化

									'''

grad_fn=<CloneBackward>，表示clone后的返回值是个中间变量，因此支持梯度的回溯。clone操作在一定程度上可以视为是一个identity-mapping函数。

detach()操作后的tensor与原始tensor共享数据内存，当原始tensor在计算图中数值发生反向传播等更新之后，detach()的tensor值也发生了改变。

注意：在pytorch中我们不要直接使用id是否相等来判断tensor是否共享内存，这只是充分条件，因为也许底层共享数据内存，但是仍然是新的tensor，比如detach()，如果我们直接打印id会出现以下情况。

				?

									import torch as t

									a = t.tensor([1.0,2.0], requires_grad=True)

									b = a.detach()

									#c[:] = a.detach()

									print(id(a))

									print(id(b))

									#140568935450520

									140570337203616

显然直接打印出来的id不等，我们可以通过简单的赋值后观察数据变化进行判断。

(2). clone()的梯度回传

detach()函数可以返回一个完全相同的tensor,与旧的tensor共享内存，脱离计算图，不会牵扯梯度计算。

而clone充当中间变量，会将梯度传给源张量进行叠加，但是本身不保存其grad，即值为None

				?

									import torch

									a = torch.tensor(1.0, requires_grad=True)

									a_ = a.clone()

									y = a**2

									z = a ** 2+a_ * 3

									y.backward()

									print(a.grad) # 2

									z.backward()

									print(a_.grad)　　　# None. 中间variable，无grad

									print(a.grad) 

									'''

									输出：

									tensor(2.) 

									None

									tensor(7.) # 2*2+3=7

									'''

使用torch.clone()获得的新tensor和原来的数据不再共享内存，但仍保留在计算图中，clone操作在不共享数据内存的同时支持梯度梯度传递与叠加，所以常用在神经网络中某个单元需要重复使用的场景下。

通常如果原tensor的requires_grad=True，则：

clone()操作后的tensor requires_grad=True
detach()操作后的tensor requires_grad=False。

				?

									import torch

									torch.manual_seed(0)

									x= torch.tensor([1., 2.], requires_grad=True)

									clone_x = x.clone() 

									detach_x = x.detach()

									clone_detach_x = x.clone().detach() 

									f = torch.nn.Linear(2, 1)

									y = f(x)

									y.backward()

									print(x.grad)

									print(clone_x.requires_grad)

									print(clone_x.grad)

									print(detach_x.requires_grad)

									print(clone_detach_x.requires_grad)

									'''

									输出结果如下：

									tensor([-0.0053, 0.3793])

									True

									None

									False

									False

									'''

另一个比较特殊的是当源张量的 require_grad=False，clone后的张量 require_grad=True，此时不存在张量回传现象，可以得到clone后的张量求导。

如下：

				?

									import torch

									a = torch.tensor(1.0)

									a_ = a.clone()

									a_.requires_grad_() #require_grad=True

									y = a_ ** 2

									y.backward()

									print(a.grad) # None

									print(a_.grad) 

									'''

									输出：

									None

									tensor(2.)

									'''