PyTorch中的拷貝和就地操作總結

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.)
'''

總結：

torch.detach() —新的tensor會脫離計算圖，不會牽扯梯度計算
torch.clone() — 新的tensor充當中間變量，會保留在計算圖中，參與梯度計算（回傳疊加），但是一般不會保留自身梯度。
原地操作(in-place, such as resize_ / resize_as_ / set_ / transpose_) 在上面兩者中執行都會引發錯誤或者警告。

引用官方文檔的話：如果你使用了in-place operation而沒有報錯的話，那麼你可以確定你的梯度計算是正確的。另外儘量避免in-place的使用。