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

torch.cuda 用于设置和运行 CUDA 操作。 它跟踪当前选择的 GPU，默认情况下，您分配的所有 CUDA 张量将在该设备上创建。 可以使用 torch.cuda.device 上下文管理器更改所选设备。

但是，一旦分配了张量，就可以对它进行操作，而与所选设备无关，并且结果将始终与张量放在同一设备上。

默认情况下，除 copy_() 和其他具有类似复制功能的方法(例如 to() 和 cuda() 。 除非您启用对等内存访问，否则尝试在分布在不同设备上的张量上启动 ops 都会引发错误。

在下面，您可以找到一个小例子，展示了这一点：

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)

异步执行

默认情况下，GPU 操作是异步的。 当您调用使用 GPU 的函数时，的操作会排队到特定的设备，但不一定要等到以后执行。 这使我们能够并行执行更多计算，包括在 CPU 或其他 GPU 上的操作。

通常，调用者看不到异步计算的效果，因为(1）每个设备按照它们排队的顺序执行操作，并且(2）当在 CPU 和 GPU 之间或两个 GPU 之间复制数据时，PyTorch 自动执行必要的同步。 因此，计算将像每个操作都同步执行一样进行。

您可以通过设置环境变量CUDA_LAUNCH_BLOCKING=1来强制进行同步计算。 当 GPU 上发生错误时，这很方便。 (对于异步执行，直到实际执行该操作后才报告这种错误，因此堆栈跟踪不会显示请求的位置。）

异步计算的结果是没有同步的时间测量不准确。 要获得精确的测量结果，应在测量之前致电 torch.cuda.synchronize() ，或使用 torch.cuda.Event 记录时间，如下所示：

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_() 等几个函数都允许使用显式non_blocking参数，该参数使调用者在不需要时绕过同步。 另一个例外是 CUDA 流，如下所述。

CUDA 流

CUDA 流是属于特定设备的线性执行序列。 通常，您无需显式创建一个：默认情况下，每个设备使用其自己的“默认”流。

每个流内部的操作都按照创建顺序进行序列化，但是来自不同流的操作可以以任何相对顺序并发执行，除非显式同步功能(例如 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)

如上所述，当“当前流”为默认流时，PyTorch 会在数据四处移动时自动执行必要的同步。 但是，使用非默认流时，用户有责任确保正确的同步。

内存管理

PyTorch 使用缓存内存分配器来加速内存分配。 这允许快速的内存重新分配而无需设备同步。 但是，分配器管理的未使用内存仍将显示为nvidia-smi中使用的内存。 您可以使用 memory_allocated() 和 max_memory_allocated() 来监视张量占用的内存，并使用  memory_reserved()  和 max_memory_reserved() 监视由缓存分配器管理的内存总量。 调用 empty_cache() 会从 PyTorch 释放所有未使用的缓存内存，以便其他 GPU 应用程序可以使用它们。 但是，张量占用的 GPU 内存不会被释放，因此不会增加可用于 PyTorch 的 GPU 内存量。

对于更高级的用户，我们通过 memory_stats() 提供更全面的内存基准测试。 我们还提供了通过 memory_snapshot() 捕获内存分配器状态的完整快照的功能，它可以帮助您了解代码所产生的基础分配模式。

cuFFT 计划缓存

对于每个 CUDA 设备，使用 cuFFT 计划的 LRU 缓存来加速在具有相同配置的相同几何形状的 CUDA 张量上重复运行 FFT 方法(例如 torch.fft())。 由于某些 cuFFT 计划可能会分配 GPU 内存，因此这些缓存具有最大容量。

您可以使用以下 API 控制和查询当前设备的缓存的属性：

• torch.backends.cuda.cufft_plan_cache.max_size给出了缓存的容量(在 CUDA 10 及更高版本上，默认值为 4096；在较旧 CUDA 版本上，默认值为 1023）。 设置该值将直接修改容量。
• torch.backends.cuda.cufft_plan_cache.size给出当前驻留在缓存中的计划数量。
• torch.backends.cuda.cufft_plan_cache.clear()清除缓存。

要控制和查询非默认设备的计划缓存，您可以使用torch.device对象或设备索引为torch.backends.cuda.cufft_plan_cache对象建立索引，并访问上述属性之一。 例如，要设置设备1的缓存容量，可以写入torch.backends.cuda.cufft_plan_cache[1].max_size = 10。

最佳实务

与设备无关的代码

由于 PyTorch 的结构，您可能需要显式编写与设备无关的代码(CPU 或 GPU）； 一个例子可能是创建一个新的张量作为循环神经网络的初始隐藏状态。

第一步是确定是否应使用 GPU。 一种常见的模式是与 is_available() 结合使用 Python 的argparse模块读取用户参数，并具有可用于禁用 CUDA 的标志。 在下面，args.device产生一个torch.device对象，该对象可用于将张量移动到 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')

现在我们有了args.device，我们可以使用它在所需设备上创建张量。

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

在许多情况下可以使用它来生成设备不可知代码。 以下是使用数据加载器时的示例：

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

在系统上使用多个 GPU 时，可以使用CUDA_VISIBLE_DEVICES环境标志来管理 PyTorch 可以使用哪些 GPU。 如上所述，要手动控制在哪个 GPU 上创建张量，最佳实践是使用 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

如果您具有张量，并且想在同一设备上创建相同类型的新张量，则可以使用torch.Tensor.new_*方法(请参见 torch.Tensor)。 前面提到的torch.*工厂函数 (Creation Ops)取决于当前 GPU 上下文和您传入的属性参数，torch.Tensor.new_*方法保留设备和张量的其他属性。

这是在创建模块时的推荐做法，在这些模块中，在前向传递期间需要在内部创建新的张量。

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]])

如果要创建与其他张量相同类型和大小的张量，并用一个或零填充，请提供 ones_like() 或 zeros_like() 作为方便的助手 函数(还保留张量的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)

使用固定的内存缓冲区

主机到 GPU 副本源自固定(页面锁定）内存时，速度要快得多。 CPU 张量和存储公开了 pin_memory() 方法，该方法返回对象的副本，并将数据放在固定的区域中。

此外，一旦固定张量或存储，就可以使用异步 GPU 副本。 只需将附加的non_blocking=True参数传递给 to() 或 >cuda() 调用。 这可用于将数据传输与计算重叠。

通过将pin_memory=True传递给其构造函数，可以使 DataLoader 返回放置在固定内存中的批处理。

使用 nn.DataParallel 代替并行处理

大多数涉及批处理输入和多个 GPU 的用例应默认使用 DataParallel 来利用多个 GPU。 即使使用 GIL，单个 Python 进程也可以使多个 GPU 饱和。

从 0.1.9 版开始，可能无法充分利用大量 GPU(8+）。 但是，这是一个正在积极开发的已知问题。 与往常一样，测试您的用例。

使用 multiprocessing 的 CUDA 模型有很多警告； 除非注意要完全满足数据处理要求，否则您的程序可能会出现错误或不确定的行为。