PyTorch CUDA 語(yǔ)義

原文： https://pytorch.org/docs/stable/notes/cuda.html

torch.cuda 用于設(shè)置和運(yùn)行 CUDA 操作。 它跟蹤當(dāng)前選擇的 GPU，默認(rèn)情況下，您分配的所有 CUDA 張量將在該設(shè)備上創(chuàng)建。 可以使用 torch.cuda.device 上下文管理器更改所選設(shè)備。

但是，一旦分配了張量，就可以對(duì)它進(jìn)行操作，而與所選設(shè)備無(wú)關(guān)，并且結(jié)果將始終與張量放在同一設(shè)備上。

默認(rèn)情況下，除 copy_() 和其他具有類似復(fù)制功能的方法(例如 to() 和 cuda() 。 除非您啟用對(duì)等內(nèi)存訪問(wèn)，否則嘗試在分布在不同設(shè)備上的張量上啟動(dòng) ops 都會(huì)引發(fā)錯(cuò)誤。

在下面，您可以找到一個(gè)小例子，展示了這一點(diǎn)：

cuda = torch.device('cuda')     # Default CUDA device
cuda0 = torch.device('cuda:0')
cuda2 = torch.device('cuda:2')  # GPU 2 (these are 0-indexed)


x = torch.tensor([1., 2.], device=cuda0)
## x.device is device(type='cuda', index=0)
y = torch.tensor([1., 2.]).cuda()
## y.device is device(type='cuda', index=0)


with torch.cuda.device(1):
    # allocates a tensor on GPU 1
    a = torch.tensor([1., 2.], device=cuda)


    # transfers a tensor from CPU to GPU 1
    b = torch.tensor([1., 2.]).cuda()
    # a.device and b.device are device(type='cuda', index=1)


    # You can also use ``Tensor.to`` to transfer a tensor:
    b2 = torch.tensor([1., 2.]).to(device=cuda)
    # b.device and b2.device are device(type='cuda', index=1)


    c = a + b
    # c.device is device(type='cuda', index=1)


    z = x + y
    # z.device is device(type='cuda', index=0)


    # even within a context, you can specify the device
    # (or give a GPU index to the .cuda call)
    d = torch.randn(2, device=cuda2)
    e = torch.randn(2).to(cuda2)
    f = torch.randn(2).cuda(cuda2)
    # d.device, e.device, and f.device are all device(type='cuda', index=2)

異步執(zhí)行

默認(rèn)情況下，GPU 操作是異步的。 當(dāng)您調(diào)用使用 GPU 的函數(shù)時(shí)，的操作會(huì)排隊(duì)到特定的設(shè)備，但不一定要等到以后執(zhí)行。 這使我們能夠并行執(zhí)行更多計(jì)算，包括在 CPU 或其他 GPU 上的操作。

通常，調(diào)用者看不到異步計(jì)算的效果，因?yàn)?1）每個(gè)設(shè)備按照它們排隊(duì)的順序執(zhí)行操作，并且(2）當(dāng)在 CPU 和 GPU 之間或兩個(gè) GPU 之間復(fù)制數(shù)據(jù)時(shí)，PyTorch 自動(dòng)執(zhí)行必要的同步。 因此，計(jì)算將像每個(gè)操作都同步執(zhí)行一樣進(jìn)行。

您可以通過(guò)設(shè)置環(huán)境變量CUDA_LAUNCH_BLOCKING=1來(lái)強(qiáng)制進(jìn)行同步計(jì)算。 當(dāng) GPU 上發(fā)生錯(cuò)誤時(shí)，這很方便。 (對(duì)于異步執(zhí)行，直到實(shí)際執(zhí)行該操作后才報(bào)告這種錯(cuò)誤，因此堆棧跟蹤不會(huì)顯示請(qǐng)求的位置。）

異步計(jì)算的結(jié)果是沒(méi)有同步的時(shí)間測(cè)量不準(zhǔn)確。 要獲得精確的測(cè)量結(jié)果，應(yīng)在測(cè)量之前致電 torch.cuda.synchronize() ，或使用 torch.cuda.Event 記錄時(shí)間，如下所示：

start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()


## Run some things here


end_event.record()
torch.cuda.synchronize()  # Wait for the events to be recorded!
elapsed_time_ms = start_event.elapsed_time(end_event)

作為例外， to() 和 copy_() 等幾個(gè)函數(shù)都允許使用顯式non_blocking參數(shù)，該參數(shù)使調(diào)用者在不需要時(shí)繞過(guò)同步。 另一個(gè)例外是 CUDA 流，如下所述。

CUDA 流

CUDA 流是屬于特定設(shè)備的線性執(zhí)行序列。 通常，您無(wú)需顯式創(chuàng)建一個(gè)：默認(rèn)情況下，每個(gè)設(shè)備使用其自己的“默認(rèn)”流。

每個(gè)流內(nèi)部的操作都按照創(chuàng)建順序進(jìn)行序列化，但是來(lái)自不同流的操作可以以任何相對(duì)順序并發(fā)執(zhí)行，除非顯式同步功能(例如 synchronize() 或 wait_stream())。 例如，以下代碼不正確：

cuda = torch.device('cuda')
s = torch.cuda.Stream()  # Create a new stream.
A = torch.empty((100, 100), device=cuda).normal_(0.0, 1.0)
with torch.cuda.stream(s):
    # sum() may start execution before normal_() finishes!
    B = torch.sum(A)

如上所述，當(dāng)“當(dāng)前流”為默認(rèn)流時(shí)，PyTorch 會(huì)在數(shù)據(jù)四處移動(dòng)時(shí)自動(dòng)執(zhí)行必要的同步。 但是，使用非默認(rèn)流時(shí)，用戶有責(zé)任確保正確的同步。

內(nèi)存管理

PyTorch 使用緩存內(nèi)存分配器來(lái)加速內(nèi)存分配。 這允許快速的內(nèi)存重新分配而無(wú)需設(shè)備同步。 但是，分配器管理的未使用內(nèi)存仍將顯示為nvidia-smi中使用的內(nèi)存。 您可以使用 memory_allocated() 和 max_memory_allocated() 來(lái)監(jiān)視張量占用的內(nèi)存，并使用  memory_reserved()  和 max_memory_reserved() 監(jiān)視由緩存分配器管理的內(nèi)存總量。 調(diào)用 empty_cache() 會(huì)從 PyTorch 釋放所有未使用的緩存內(nèi)存，以便其他 GPU 應(yīng)用程序可以使用它們。 但是，張量占用的 GPU 內(nèi)存不會(huì)被釋放，因此不會(huì)增加可用于 PyTorch 的 GPU 內(nèi)存量。

對(duì)于更高級(jí)的用戶，我們通過(guò) memory_stats() 提供更全面的內(nèi)存基準(zhǔn)測(cè)試。 我們還提供了通過(guò) memory_snapshot() 捕獲內(nèi)存分配器狀態(tài)的完整快照的功能，它可以幫助您了解代碼所產(chǎn)生的基礎(chǔ)分配模式。

cuFFT 計(jì)劃緩存

對(duì)于每個(gè) CUDA 設(shè)備，使用 cuFFT 計(jì)劃的 LRU 緩存來(lái)加速在具有相同配置的相同幾何形狀的 CUDA 張量上重復(fù)運(yùn)行 FFT 方法(例如 torch.fft())。 由于某些 cuFFT 計(jì)劃可能會(huì)分配 GPU 內(nèi)存，因此這些緩存具有最大容量。

您可以使用以下 API 控制和查詢當(dāng)前設(shè)備的緩存的屬性：

• torch.backends.cuda.cufft_plan_cache.max_size給出了緩存的容量(在 CUDA 10 及更高版本上，默認(rèn)值為 4096；在較舊 CUDA 版本上，默認(rèn)值為 1023）。 設(shè)置該值將直接修改容量。
• torch.backends.cuda.cufft_plan_cache.size給出當(dāng)前駐留在緩存中的計(jì)劃數(shù)量。
• torch.backends.cuda.cufft_plan_cache.clear()清除緩存。

要控制和查詢非默認(rèn)設(shè)備的計(jì)劃緩存，您可以使用torch.device對(duì)象或設(shè)備索引為torch.backends.cuda.cufft_plan_cache對(duì)象建立索引，并訪問(wèn)上述屬性之一。 例如，要設(shè)置設(shè)備1的緩存容量，可以寫入torch.backends.cuda.cufft_plan_cache[1].max_size = 10。

最佳實(shí)務(wù)

與設(shè)備無(wú)關(guān)的代碼

由于 PyTorch 的結(jié)構(gòu)，您可能需要顯式編寫與設(shè)備無(wú)關(guān)的代碼(CPU 或 GPU）； 一個(gè)例子可能是創(chuàng)建一個(gè)新的張量作為循環(huán)神經(jīng)網(wǎng)絡(luò)的初始隱藏狀態(tài)。

第一步是確定是否應(yīng)使用 GPU。 一種常見(jiàn)的模式是與 is_available() 結(jié)合使用 Python 的argparse模塊讀取用戶參數(shù)，并具有可用于禁用 CUDA 的標(biāo)志。 在下面，args.device產(chǎn)生一個(gè)torch.device對(duì)象，該對(duì)象可用于將張量移動(dòng)到 CPU 或 CUDA。

import argparse
import torch


parser = argparse.ArgumentParser(description='PyTorch Example')
parser.add_argument('--disable-cuda', action='store_true',
                    help='Disable CUDA')
args = parser.parse_args()
args.device = None
if not args.disable_cuda and torch.cuda.is_available():
    args.device = torch.device('cuda')
else:
    args.device = torch.device('cpu')

現(xiàn)在我們有了args.device，我們可以使用它在所需設(shè)備上創(chuàng)建張量。

x = torch.empty((8, 42), device=args.device)
net = Network().to(device=args.device)

在許多情況下可以使用它來(lái)生成設(shè)備不可知代碼。 以下是使用數(shù)據(jù)加載器時(shí)的示例：

cuda0 = torch.device('cuda:0')  # CUDA GPU 0
for i, x in enumerate(train_loader):
    x = x.to(cuda0)

在系統(tǒng)上使用多個(gè) GPU 時(shí)，可以使用CUDA_VISIBLE_DEVICES環(huán)境標(biāo)志來(lái)管理 PyTorch 可以使用哪些 GPU。 如上所述，要手動(dòng)控制在哪個(gè) GPU 上創(chuàng)建張量，最佳實(shí)踐是使用 torch.cuda.device 上下文管理器。

print("Outside device is 0")  # On device 0 (default in most scenarios)
with torch.cuda.device(1):
    print("Inside device is 1")  # On device 1
print("Outside device is still 0")  # On device 0

如果您具有張量，并且想在同一設(shè)備上創(chuàng)建相同類型的新張量，則可以使用torch.Tensor.new_*方法(請(qǐng)參見(jiàn) torch.Tensor)。 前面提到的torch.*工廠函數(shù) (Creation Ops)取決于當(dāng)前 GPU 上下文和您傳入的屬性參數(shù)，torch.Tensor.new_*方法保留設(shè)備和張量的其他屬性。

這是在創(chuàng)建模塊時(shí)的推薦做法，在這些模塊中，在前向傳遞期間需要在內(nèi)部創(chuàng)建新的張量。

cuda = torch.device('cuda')
x_cpu = torch.empty(2)
x_gpu = torch.empty(2, device=cuda)
x_cpu_long = torch.empty(2, dtype=torch.int64)


y_cpu = x_cpu.new_full([3, 2], fill_value=0.3)
print(y_cpu)


    tensor([[ 0.3000,  0.3000],
            [ 0.3000,  0.3000],
            [ 0.3000,  0.3000]])


y_gpu = x_gpu.new_full([3, 2], fill_value=-5)
print(y_gpu)


    tensor([[-5.0000, -5.0000],
            [-5.0000, -5.0000],
            [-5.0000, -5.0000]], device='cuda:0')


y_cpu_long = x_cpu_long.new_tensor([[1, 2, 3]])
print(y_cpu_long)


    tensor([[ 1,  2,  3]])

如果要?jiǎng)?chuàng)建與其他張量相同類型和大小的張量，并用一個(gè)或零填充，請(qǐng)?zhí)峁?nbsp;ones_like() 或 zeros_like() 作為方便的助手 函數(shù)(還保留張量的torch.device和torch.dtype）。

x_cpu = torch.empty(2, 3)
x_gpu = torch.empty(2, 3)


y_cpu = torch.ones_like(x_cpu)
y_gpu = torch.zeros_like(x_gpu)

使用固定的內(nèi)存緩沖區(qū)

主機(jī)到 GPU 副本源自固定(頁(yè)面鎖定）內(nèi)存時(shí)，速度要快得多。 CPU 張量和存儲(chǔ)公開(kāi)了 pin_memory() 方法，該方法返回對(duì)象的副本，并將數(shù)據(jù)放在固定的區(qū)域中。

此外，一旦固定張量或存儲(chǔ)，就可以使用異步 GPU 副本。 只需將附加的non_blocking=True參數(shù)傳遞給 to() 或 >cuda() 調(diào)用。 這可用于將數(shù)據(jù)傳輸與計(jì)算重疊。

通過(guò)將pin_memory=True傳遞給其構(gòu)造函數(shù)，可以使 DataLoader 返回放置在固定內(nèi)存中的批處理。

使用 nn.DataParallel 代替并行處理

大多數(shù)涉及批處理輸入和多個(gè) GPU 的用例應(yīng)默認(rèn)使用 DataParallel 來(lái)利用多個(gè) GPU。 即使使用 GIL，單個(gè) Python 進(jìn)程也可以使多個(gè) GPU 飽和。

從 0.1.9 版開(kāi)始，可能無(wú)法充分利用大量 GPU(8+）。 但是，這是一個(gè)正在積極開(kāi)發(fā)的已知問(wèn)題。 與往常一樣，測(cè)試您的用例。

使用 multiprocessing 的 CUDA 模型有很多警告； 除非注意要完全滿足數(shù)據(jù)處理要求，否則您的程序可能會(huì)出現(xiàn)錯(cuò)誤或不確定的行為。