PyTorch中的In-place操作是什麼?為什麼要避免使用這種操作?

語言: CN / TW / HK

持續創作,加速成長!這是我參與「掘金日新計劃 · 10 月更文挑戰」的第4天,點選檢視活動詳情

導讀

In-place操作用在推理的時候可以顯著節省記憶體,但是訓練的時候一定要小心使用。

圖片

如今的高階深度神經網路擁有數百萬個可訓練引數,訓練它們通常會導致GPU記憶體耗盡。有幾種簡單的方法可以減少模型佔用的GPU記憶體,例如:

  • 考慮改變模型的架構或使用具有較少可訓練引數的模型型別(例如,選擇DenseNet-121而不是DenseNet-169)。這種方法會影響模型的效能度量。
  • 減少batch大小或手動設定資料載入workers的數量。在這種情況下,模型需要更長的時間來訓練。

在神經網路中使用in-place操作可能有助於避免上述方法的缺點,同時節省一些GPU記憶體。但是,由於幾個原因,不建議使用in-place操作。

在這篇文章中,內容包括:

  • 描述什麼是in-place操作,並演示他們如何可能有助於節省GPU記憶體。
  • 告訴我們為什麼要避免in-place操作或非常小心地使用它們。

In-place 操作

“In-place運算是一種直接改變給定線性函式、向量、矩陣(張量)內容而不復制的運算。"

根據定義,in-place操作不會複製輸入。這就是為什麼它們可以幫助在操作高維資料時減少記憶體使用。

我想演示in-place操作如何幫助消耗更少的GPU記憶體。為了做到這一點,我將度量out- place ReLU和in-place ReLU分配的記憶體,使用這個簡單的函式:

```

Import PyTorch

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() / 10242     start_memory = torch.cuda.memory_allocated() / 10242          # 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() / 10242     end_memory = torch.cuda.memory_allocated() / 10242          # Return amount of memory allocated for ReLU call     return end_memory - start_memory, end_max_memory - start_max_memory ```

呼叫該函式來測量out- place ReLU分配的記憶體:

```

setup the device

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

call the function to measure allocated memory

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

然後呼叫in-place ReLU:

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

看來使用in-place操作可以幫助我們節省一些GPU記憶體。然而,我們在使用現場操作時應該非常謹慎,並且要反覆檢查。在接下來的部分,我將告訴你為什麼。

In-place 操作的缺點

in-place操作的主要缺點是,它們可能會覆蓋計算梯度所需的值,這意味著破壞模型的訓練過程。這是PyTorch autograd官方文件所說的:

在autograd支援in-place操作是一件困難的事情,我們在大多數情況下不鼓勵使用它們。Autograd的主動緩衝區釋放和重用使其非常高效,在很少情況下,in-place操作實際上會顯著降低記憶體使用量。除非你正在承受巨大的記憶體壓力,否則你可能永遠都不需要使用它們。

限制in-place作業的適用性的主要原因有兩個:

1、in-place操作可能會覆蓋計算梯度所需的值。

2、每個in-place操作實際上都需要實現重寫計算圖。Out-of-place版本只是簡單地分配新物件並保持對舊圖的引用,而in-place操作則要求將所有輸入的建立者更改為表示該操作的函式。

要小心使用in-place操作的另一個原因是,它們的實現非常棘手。這就是為什麼我建議使用PyTorch標準的in-place操作(如上面的就地ReLU),而不是手動實現。

讓我們看一個SiLU(或Swish-1)啟用函式的例子。以下是SiLU的out-of-place實現:

def silu(input):     '''     Out-of-place implementation of SiLU activation function     https://arxiv.org/pdf/1606.08415.pdf     '''     return input * torch.sigmoid(input)

我們嘗試使用torch.sigmoid_實現一個 in-place SiLU:

def silu_inplace_1(input):     '''     Incorrect implementation of in-place SiLU activation function     https://arxiv.org/pdf/1606.08415.pdf     '''     return input * torch.sigmoid_(input) # THIS IS INCORRECT!!!

上面的程式碼不正確地實現了in-place SiLU。只要比較兩個函式返回的值,就可以確定。實際上,函式silu_inplace_1 返回sigmoid(input) * sigmoid(input)!使用torch.sigmoid_實現SiLU的例子:

def silu_inplace_2(input):     '''     Example of implementation of in-place SiLU activation function using torch.sigmoid_     https://arxiv.org/pdf/1606.08415.pdf     '''     result = input.clone()     torch.sigmoid_(input)     input *= result     return input

這個小示例演示了為什麼我們在使用in-place操作時應該謹慎並檢查兩次。

總結

本文簡介:

  • 說明了in-plac及其目的。演示了in-plac操作如何幫助消耗更少的GPU記憶體。
  • 描述了in-plac操作的顯著缺點。人們應該非常小心地使用它們,並檢查兩次結果。

英文原文:https://towardsdatascience.com/in-place-operations-in-pytorch-f91d493e970e