A guide on good usage of non_blocking 和 pin_memory() 在 PyTorch

创建日期: 2024年7月31日 | 最后更新日期: 2024年8月1日 | 最后验证日期: 2024年11月5日

作者: Vincent Moens

介绍

从CPU传输数据到GPU在许多PyTorch应用中都是基础操作。 用户需要了解在设备间移动数据时最有效的工具和选项。 本教程将探讨PyTorch中的两种关键方法，用于设备到设备的数据传输： pin_memory() 和 to()，带有 non_blocking=True 选项。

你将学习的内容

通过异步传输和内存绑定可以优化张量从CPU到GPU的转移，但需要注意一些重要事项：

• 使用 tensor.pin_memory().to(device, non_blocking=True) 可能比直接的 tensor.to(device) 慢两倍。

• 一般而言，tensor.to(device, non_blocking=True) 是一个有效的选择，可以提高传输速度。

• While cpu_tensor.to("cuda", non_blocking=True).mean() 执行正确，尝试执行 cuda_tensor.to("cpu", non_blocking=True).mean() 将导致错误输出。

引言

本教程中报告的性能取决于构建教程所使用的系统。 尽管结论适用于不同的系统，但具体的观察结果可能会因可用硬件的不同而略有差异，尤其是在较旧的硬件上。 本教程的主要目标是提供一个理论框架，以便理解CPU到GPU的数据传输。 然而，任何设计决策都应针对个别情况进行调整，并由基准带宽测量值指导，同时还要考虑当前任务的具体要求。

import torch

assert torch.cuda.is_available(), "A cuda device is required to run this tutorial"

本教程要求安装 tensordict。如果你的环境中还没有 tensordict，请在单独的单元格中运行以下命令进行安装：

# Install tensordict with the following command
!pip3 install tensordict

我们首先概述这些概念背后的理论，然后转向具体的功能测试示例。

背景

内存管理基础

当在Pytorch中创建一个CPU张量时，这个张量的内容需要被放置到内存中。这里的内存是一个相当复杂的概念，值得仔细查看。 我们区分两种由内存管理单元处理的内存类型：RAM（为了简化）和磁盘上的交换空间（可能是硬盘）。一起，磁盘上的可用空间和RAM（物理内存）组成了虚拟内存，这是一种对总资源进行抽象的方式。 简而言之，虚拟内存使得可用空间比单独的RAM要大，并且创造了一种错觉，即主内存比实际更大。

在正常情况下，常规CPU张量是可以分页的，这意味着它被分成称为页面的块，这些块可以在虚拟内存中的任意位置存在（既可以在RAM中也可以在磁盘上）。正如之前提到的，这有一个优势，即内存似乎比实际的主要内存更大。

通常，当程序访问未在RAM中的页面时，会发生“页面故障”（page fault），操作系统（OS）随后会将该页面加载到RAM中（“交换入”或“页面入”）。 相应地，为了给新页面腾出空间，操作系统可能需要将另一个页面交换出去（或“页面出”）。

与可分页内存不同，固定（或页面锁定或非分页）内存是一种不能被交换到磁盘的内存类型。 它允许更快更可预测的访问时间，但缺点是它的容量比可分页内存（即主内存）更为有限。

CUDA和（非）页面式内存

要理解CUDA如何将张量从CPU复制到CUDA，请考虑上述两种情况：

• 如果内存被页面锁定，设备可以直接在主内存中访问该内存，地址定义清晰，需要读取这些数据的函数可以显著加速。

• 如果内存可页面化，所有页面都必须先被加载到主内存，然后再发送到GPU。 这个操作可能需要花费一些时间，并且在执行时不如页面锁定的张量那么可预测。

更确切地说，当CUDA将CPU上的可页面数据发送到GPU时，它必须首先创建一个页面锁定的副本，然后再进行数据传输。

异步与同步操作（使用 non_blocking=True （CUDA cudaMemcpyAsync））

在从主机（例如CPU）复制数据到设备（例如GPU）时，CUDA工具包提供了同步或异步执行这些操作的方式。

在实践中，当调用to()时，PyTorch 总是会调用 cudaMemcpyAsync。 如果non_blocking=False（默认值），每次调用to()后都会触发一个同步操作，使得主线程中的调用to()变为阻塞的。 如果non_blocking=True，则不会触发同步操作，主机上的主线程也不会被阻塞。 因此，从主机的角度来看，可以同时发送多个张量到设备，因为线程不需要等待一次传输完成后再发起下一次传输。

注意

一般来说，数据传输会阻塞设备端（即使主机端没有阻塞）： 在另一个操作执行时，设备上的复制操作不能发生。 然而，在某些高级场景中，复制和内核执行可以在GPU端同时进行。 如以下示例所示，必须满足三个要求才能实现这一点：

1. 设备至少需要有一个空闲的DMA（直接内存访问）引擎。现代GPU架构，如Volterra、Tesla或H100设备，拥有多个DMA引擎。

2. 转移必须在单独的非默认cuda流上进行。在PyTorch中，cuda流可以使用 Stream 处理。

3. 源数据必须位于固定内存中。

我们通过运行以下脚本的性能分析来演示这一点。

import contextlib

from torch.cuda import Stream


s = Stream()

torch.manual_seed(42)
t1_cpu_pinned = torch.randn(1024**2 * 5, pin_memory=True)
t2_cpu_paged = torch.randn(1024**2 * 5, pin_memory=False)
t3_cuda = torch.randn(1024**2 * 5, device="cuda:0")

assert torch.cuda.is_available()
device = torch.device("cuda", torch.cuda.current_device())


# The function we want to profile
def inner(pinned: bool, streamed: bool):
    with torch.cuda.stream(s) if streamed else contextlib.nullcontext():
        if pinned:
            t1_cuda = t1_cpu_pinned.to(device, non_blocking=True)
        else:
            t2_cuda = t2_cpu_paged.to(device, non_blocking=True)
        t_star_cuda_h2d_event = s.record_event()
    # This operation can be executed during the CPU to GPU copy if and only if the tensor is pinned and the copy is
    #  done in the other stream
    t3_cuda_mul = t3_cuda * t3_cuda * t3_cuda
    t3_cuda_h2d_event = torch.cuda.current_stream().record_event()
    t_star_cuda_h2d_event.synchronize()
    t3_cuda_h2d_event.synchronize()


# Our profiler: profiles the `inner` function and stores the results in a .json file
def benchmark_with_profiler(
    pinned,
    streamed,
) -> None:
    torch._C._profiler._set_cuda_sync_enabled_val(True)
    wait, warmup, active = 1, 1, 2
    num_steps = wait + warmup + active
    rank = 0
    with torch.profiler.profile(
        activities=[
            torch.profiler.ProfilerActivity.CPU,
            torch.profiler.ProfilerActivity.CUDA,
        ],
        schedule=torch.profiler.schedule(
            wait=wait, warmup=warmup, active=active, repeat=1, skip_first=1
        ),
    ) as prof:
        for step_idx in range(1, num_steps + 1):
            inner(streamed=streamed, pinned=pinned)
            if rank is None or rank == 0:
                prof.step()
    prof.export_chrome_trace(f"trace_streamed{int(streamed)}_pinned{int(pinned)}.json")

加载这些性能踪迹到 Chrome 中（chrome://tracing）会显示以下结果：首先，让我们看看如果在主线程中将可分页张量发送到 GPU 后再执行第 t3_cuda 步的算术运算会发生什么：

benchmark_with_profiler(streamed=False, pinned=False)

固定张量不会显著改变追踪过程，两个操作仍然会依次执行：

benchmark_with_profiler(streamed=False, pinned=True)

将可分页张量发送到单独流的GPU也是一个阻塞操作：

benchmark_with_profiler(streamed=True, pinned=False)

只有固定张量在单独流上复制到GPU并与主线流上执行的另一个CUDA内核重叠：

benchmark_with_profiler(streamed=True, pinned=True)

Pytorch视角

pin_memory()

PyTorch 提供了通过 pin_memory() 方法和构造函数参数创建并发送张量到页面锁定内存的可能性。 在初始化了 CUDA 的机器上，可以通过 pin_memory() 方法将 CPU 张量转换为页面锁定内存。重要的是，pin_memory 在主机主线程上是阻塞的：它会在执行下一个操作之前等待张量被复制到页面锁定内存中。 可以直接使用类似 zeros()、ones() 和其他构造函数在页面锁定内存中创建新的张量。

让我们检查一下固定内存的速度以及将张量发送到CUDA的速度：

import torch
import gc
from torch.utils.benchmark import Timer
import matplotlib.pyplot as plt


def timer(cmd):
    median = (
        Timer(cmd, globals=globals())
        .adaptive_autorange(min_run_time=1.0, max_run_time=20.0)
        .median
        * 1000
    )
    print(f"{cmd}: {median: 4.4f} ms")
    return median


# A tensor in pageable memory
pageable_tensor = torch.randn(1_000_000)

# A tensor in page-locked (pinned) memory
pinned_tensor = torch.randn(1_000_000, pin_memory=True)

# Runtimes:
pageable_to_device = timer("pageable_tensor.to('cuda:0')")
pinned_to_device = timer("pinned_tensor.to('cuda:0')")
pin_mem = timer("pageable_tensor.pin_memory()")
pin_mem_to_device = timer("pageable_tensor.pin_memory().to('cuda:0')")

# Ratios:
r1 = pinned_to_device / pageable_to_device
r2 = pin_mem_to_device / pageable_to_device

# Create a figure with the results
fig, ax = plt.subplots()

xlabels = [0, 1, 2]
bar_labels = [
    "pageable_tensor.to(device) (1x)",
    f"pinned_tensor.to(device) ({r1:4.2f}x)",
    f"pageable_tensor.pin_memory().to(device) ({r2:4.2f}x)"
    f"\npin_memory()={100*pin_mem/pin_mem_to_device:.2f}% of runtime.",
]
values = [pageable_to_device, pinned_to_device, pin_mem_to_device]
colors = ["tab:blue", "tab:red", "tab:orange"]
ax.bar(xlabels, values, label=bar_labels, color=colors)

ax.set_ylabel("Runtime (ms)")
ax.set_title("Device casting runtime (pin-memory)")
ax.set_xticks([])
ax.legend()

plt.show()

# Clear tensors
del pageable_tensor, pinned_tensor
_ = gc.collect()

pageable_tensor.to('cuda:0'):  0.4667 ms
pinned_tensor.to('cuda:0'):  0.3703 ms
pageable_tensor.pin_memory():  0.3607 ms
pageable_tensor.pin_memory().to('cuda:0'):  0.7238 ms

我们可以观察到，将固定内存张量转换为GPU张量确实比将可页面张量转换为GPU张量要快得多，因为实际上，可页面张量在被发送到GPU之前必须先复制到固定内存中。

然而，与一种较为常见的看法相反，在将可分页张量转换到GPU之前对它调用pin_memory()并不会带来显著的速度提升，相反，这个调用通常比直接执行数据传输要慢。这合乎情理，因为我们实际上是在要求Python执行一个CUDA会自行执行的操作，而在此之前还需要从主机复制数据到设备。

注意

The PyTorch 实现中的 pin_memory 依赖于通过 cudaHostAlloc 创建全新的固定内存存储，在某些情况下可能会比 cudaMemcpy 分块传输数据更快。 同样，观察结果也可能根据可用硬件、发送的张量大小或可用的RAM数量而有所不同。

non_blocking=True

如前所述，许多PyTorch操作可以通过non_blocking参数以异步方式在主机上执行。

在这里，为了准确计算使用 non_blocking 的好处，我们将设计一个稍微更复杂的实验，因为我们想评估在有和没有调用 non_blocking 的情况下，将多个张量发送到GPU的速度。

# A simple loop that copies all tensors to cuda
def copy_to_device(*tensors):
    result = []
    for tensor in tensors:
        result.append(tensor.to("cuda:0"))
    return result


# A loop that copies all tensors to cuda asynchronously
def copy_to_device_nonblocking(*tensors):
    result = []
    for tensor in tensors:
        result.append(tensor.to("cuda:0", non_blocking=True))
    # We need to synchronize
    torch.cuda.synchronize()
    return result


# Create a list of tensors
tensors = [torch.randn(1000) for _ in range(1000)]
to_device = timer("copy_to_device(*tensors)")
to_device_nonblocking = timer("copy_to_device_nonblocking(*tensors)")

# Ratio
r1 = to_device_nonblocking / to_device

# Plot the results
fig, ax = plt.subplots()

xlabels = [0, 1]
bar_labels = [f"to(device) (1x)", f"to(device, non_blocking=True) ({r1:4.2f}x)"]
colors = ["tab:blue", "tab:red"]
values = [to_device, to_device_nonblocking]

ax.bar(xlabels, values, label=bar_labels, color=colors)

ax.set_ylabel("Runtime (ms)")
ax.set_title("Device casting runtime (non-blocking)")
ax.set_xticks([])
ax.legend()

plt.show()

copy_to_device(*tensors):  25.1154 ms
copy_to_device_nonblocking(*tensors):  19.2452 ms

为了更好地理解这里发生了什么，让我们对这两个函数进行性能分析：

from torch.profiler import profile, ProfilerActivity


def profile_mem(cmd):
    with profile(activities=[ProfilerActivity.CPU]) as prof:
        exec(cmd)
    print(cmd)
    print(prof.key_averages().table(row_limit=10))

让我们先看看常规的to(device)调用堆栈：

print("Call to `to(device)`", profile_mem("copy_to_device(*tensors)"))

copy_to_device(*tensors)
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                     Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                 aten::to         3.75%       1.167ms       100.00%      31.089ms      31.089us          1000
           aten::_to_copy        13.71%       4.262ms        96.25%      29.922ms      29.922us          1000
      aten::empty_strided        25.21%       7.838ms        25.21%       7.838ms       7.838us          1000
              aten::copy_        19.17%       5.960ms        57.33%      17.822ms      17.822us          1000
          cudaMemcpyAsync        19.18%       5.962ms        19.18%       5.962ms       5.962us          1000
    cudaStreamSynchronize        18.98%       5.901ms        18.98%       5.901ms       5.901us          1000
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 31.089ms

Call to `to(device)` None

and now the non_blocking版本：

print(
    "Call to `to(device, non_blocking=True)`",
    profile_mem("copy_to_device_nonblocking(*tensors)"),
)

copy_to_device_nonblocking(*tensors)
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                     Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                 aten::to         4.47%       1.023ms        99.90%      22.875ms      22.875us          1000
           aten::_to_copy        16.00%       3.664ms        95.43%      21.851ms      21.851us          1000
      aten::empty_strided        32.26%       7.386ms        32.26%       7.386ms       7.386us          1000
              aten::copy_        22.86%       5.235ms        47.17%      10.801ms      10.801us          1000
          cudaMemcpyAsync        24.31%       5.566ms        24.31%       5.566ms       5.566us          1000
    cudaDeviceSynchronize         0.10%      23.144us         0.10%      23.144us      23.144us             1
-------------------------  ------------  ------------  ------------  ------------  ------------  ------------
Self CPU time total: 22.898ms

Call to `to(device, non_blocking=True)` None

The results are without any doubt better when using non_blocking=True, as all transfers are initiated simultaneously on the host side and only one synchronization is done.

根据张量的数量和大小以及使用的硬件不同，带来的好处也会有所不同。

注意

有趣的是，阻塞to("cuda")实际上执行了与non_blocking=True相同的异步设备转换操作 （cudaMemcpyAsync），只不过在每次复制后都有一个同步点。

协同效应

现在我们已经指出，将已存在于固定内存中的张量传输到GPU比从可页面内存传输要快，并且我们知道这些传输异步进行也会比同步进行更快，我们可以对这些方法的组合进行基准测试。首先，让我们编写几个新函数，在每个张量上调用pin_memory和to(device)：

def pin_copy_to_device(*tensors):
    result = []
    for tensor in tensors:
        result.append(tensor.pin_memory().to("cuda:0"))
    return result


def pin_copy_to_device_nonblocking(*tensors):
    result = []
    for tensor in tensors:
        result.append(tensor.pin_memory().to("cuda:0", non_blocking=True))
    # We need to synchronize
    torch.cuda.synchronize()
    return result

使用 pin_memory() 的好处在较大批次的大型张量中更为明显：

tensors = [torch.randn(1_000_000) for _ in range(1000)]
page_copy = timer("copy_to_device(*tensors)")
page_copy_nb = timer("copy_to_device_nonblocking(*tensors)")

tensors_pinned = [torch.randn(1_000_000, pin_memory=True) for _ in range(1000)]
pinned_copy = timer("copy_to_device(*tensors_pinned)")
pinned_copy_nb = timer("copy_to_device_nonblocking(*tensors_pinned)")

pin_and_copy = timer("pin_copy_to_device(*tensors)")
pin_and_copy_nb = timer("pin_copy_to_device_nonblocking(*tensors)")

# Plot
strategies = ("pageable copy", "pinned copy", "pin and copy")
blocking = {
    "blocking": [page_copy, pinned_copy, pin_and_copy],
    "non-blocking": [page_copy_nb, pinned_copy_nb, pin_and_copy_nb],
}

x = torch.arange(3)
width = 0.25
multiplier = 0


fig, ax = plt.subplots(layout="constrained")

for attribute, runtimes in blocking.items():
    offset = width * multiplier
    rects = ax.bar(x + offset, runtimes, width, label=attribute)
    ax.bar_label(rects, padding=3, fmt="%.2f")
    multiplier += 1

# Add some text for labels, title and custom x-axis tick labels, etc.
ax.set_ylabel("Runtime (ms)")
ax.set_title("Runtime (pin-mem and non-blocking)")
ax.set_xticks([0, 1, 2])
ax.set_xticklabels(strategies)
plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
ax.legend(loc="upper left", ncols=3)

plt.show()

del tensors, tensors_pinned
_ = gc.collect()

copy_to_device(*tensors):  616.7664 ms
copy_to_device_nonblocking(*tensors):  544.0807 ms
copy_to_device(*tensors_pinned):  371.0937 ms
copy_to_device_nonblocking(*tensors_pinned):  345.8835 ms
pin_copy_to_device(*tensors):  996.5177 ms
pin_copy_to_device_nonblocking(*tensors):  677.9741 ms

其他复制方向（GPU -> CPU，CPU -> MPS）

直到现在，我们一直假设从CPU到GPU的异步复制是安全的。 这通常是真的，因为CUDA会自动处理同步以确保读取数据时访问的数据是有效的。 然而，这种保证并不适用于相反方向的传输，即从GPU到CPU的传输。 如果没有显式的同步，这些传输无法保证在数据访问时复制已完成。因此，主机上的数据可能是不完整的或错误的，实际上变成了垃圾数据。

tensor = (
    torch.arange(1, 1_000_000, dtype=torch.double, device="cuda")
    .expand(100, 999999)
    .clone()
)
torch.testing.assert_close(
    tensor.mean(), torch.tensor(500_000, dtype=torch.double, device="cuda")
), tensor.mean()
try:
    i = -1
    for i in range(100):
        cpu_tensor = tensor.to("cpu", non_blocking=True)
        torch.testing.assert_close(
            cpu_tensor.mean(), torch.tensor(500_000, dtype=torch.double)
        )
    print("No test failed with non_blocking")
except AssertionError:
    print(f"{i}th test failed with non_blocking. Skipping remaining tests")
try:
    i = -1
    for i in range(100):
        cpu_tensor = tensor.to("cpu", non_blocking=True)
        torch.cuda.synchronize()
        torch.testing.assert_close(
            cpu_tensor.mean(), torch.tensor(500_000, dtype=torch.double)
        )
    print("No test failed with synchronize")
except AssertionError:
    print(f"One test failed with synchronize: {i}th assertion!")

0th test failed with non_blocking. Skipping remaining tests
No test failed with synchronize

同样的考虑也适用于从CPU复制到非CUDA设备，例如MPS。 通常，在目标为CUDA启用设备的情况下，异步设备复制不需要显式的同步操作即可保证安全。

总结来说，使用 non_blocking=True 从CPU复制数据到GPU是安全的，但对于其他方向， 使用 non_blocking=True 仍然可以，但用户必须确保在访问数据之前执行设备同步。

实用建议

我们现在可以根据我们的观察提出一些早期建议：

一般情况下，non_blocking=True会提供良好的吞吐量，无论原始张量是否在固定内存中。 如果张量已经在固定内存中，传输可以加速，但手动从Python主线程将它发送到固定内存是一个阻塞操作，因此会大大抵消使用non_blocking=True的优势（因为CUDA会自行处理pin_memory的传输）。

有人现在可能会合理地问，pin_memory() 方法有什么用。 在接下来的部分中，我们将进一步探讨如何使用它来更加快速地传输数据。

额外考虑事项

PyTorch notoriously 提供了一个 DataLoader 类，其构造函数接受一个 pin_memory 参数。 考虑到我们之前关于 pin_memory 的讨论，你可能会好奇 DataLoader 如何能够加速数据传输，如果内存绑定本质上是阻塞的话。

关键在于DataLoader使用了一个单独的线程来处理数据从 pageable 内存到 pinned 内存的传输，从而防止主线程出现阻塞。

为了说明这一点，我们将使用同名库中的TensorDict原语。 调用to()时，默认行为是异步将张量发送到设备，随后紧接着一次对torch.device.synchronize()的调用。

此外，TensorDict.to() 包含一个 non_blocking_pin 选项，该选项启动多个线程来执行 pin_memory()，然后再继续执行 to(device)。 这种方法可以进一步加速数据传输，如下例所示。

from tensordict import TensorDict
import torch
from torch.utils.benchmark import Timer
import matplotlib.pyplot as plt

# Create the dataset
td = TensorDict({str(i): torch.randn(1_000_000) for i in range(1000)})

# Runtimes
copy_blocking = timer("td.to('cuda:0', non_blocking=False)")
copy_non_blocking = timer("td.to('cuda:0')")
copy_pin_nb = timer("td.to('cuda:0', non_blocking_pin=True, num_threads=0)")
copy_pin_multithread_nb = timer("td.to('cuda:0', non_blocking_pin=True, num_threads=4)")

# Rations
r1 = copy_non_blocking / copy_blocking
r2 = copy_pin_nb / copy_blocking
r3 = copy_pin_multithread_nb / copy_blocking

# Figure
fig, ax = plt.subplots()

xlabels = [0, 1, 2, 3]
bar_labels = [
    "Blocking copy (1x)",
    f"Non-blocking copy ({r1:4.2f}x)",
    f"Blocking pin, non-blocking copy ({r2:4.2f}x)",
    f"Non-blocking pin, non-blocking copy ({r3:4.2f}x)",
]
values = [copy_blocking, copy_non_blocking, copy_pin_nb, copy_pin_multithread_nb]
colors = ["tab:blue", "tab:red", "tab:orange", "tab:green"]

ax.bar(xlabels, values, label=bar_labels, color=colors)

ax.set_ylabel("Runtime (ms)")
ax.set_title("Device casting runtime")
ax.set_xticks([])
ax.legend()

plt.show()

td.to('cuda:0', non_blocking=False):  619.8487 ms
td.to('cuda:0'):  542.9433 ms
td.to('cuda:0', non_blocking_pin=True, num_threads=0):  658.6780 ms
td.to('cuda:0', non_blocking_pin=True, num_threads=4):  360.2737 ms

在本例中，我们将许多大型张量从CPU传输到GPU。 这种场景非常适合利用多线程pin_memory()，这可以显著提升性能。 然而，如果张量较小，多线程相关的开销可能会超过其带来的好处。 同样，如果只有少量张量，将张量绑定到单独线程的优势也会受到限制。

需要注意的是，虽然将张量从可页面内存转移到GPU之前先在固定内存中创建永久缓冲区似乎可以加快数据传输，但实际上这种策略并不一定能加速计算。固定内存中的数据复制瓶颈仍然是一个限制因素。

此外，将磁盘上的数据（无论是共享内存还是文件）传输到GPU通常需要一个中间步骤，即将数据复制到固定内存（位于RAM中）。 在这种情况下，使用non_blocking进行大规模数据传输会显著增加RAM消耗，可能会导致不良影响。

在实践中，并没有一劳永逸的解决方案。 使用多线程 pin_memory 结合 non_blocking 转移的效果取决于多种因素，包括具体的系统、操作系统、硬件以及执行的任务的性质。 当尝试加快CPU和GPU之间的数据传输速度或比较不同场景下的吞吐量时，请检查以下因素：

• 可用的核心数量

  有多少CPU核心可用？系统是否与其他用户或可能竞争资源的进程共享？

• 核心利用率

  CPU核心是否被其他进程大量占用？应用程序在数据传输的同时是否还会执行其他CPU密集型任务？

• 内存利用率

  当前使用了多少可页面和页锁定内存？是否有足够的空闲内存可以分配额外的固定内存而不影响系统性能？请记住，没有什么是免费的，例如pin_memory会消耗RAM并可能影响其他任务。

• CUDA 设备能力

  GPU 是否支持多个DMA引擎以实现并发数据传输？正在使用的CUDA设备的具体能力和限制是什么？

• 要发送的张量数量

  在典型的操作中会传输多少张量？

• 要发送的张量大小

  传输的张量大小是多少？几个大的张量或许多小的张量可能不会从同一个传输程序中受益。

• 系统架构

  系统的架构如何影响数据传输速度（例如，总线速度、网络延迟）？

此外，在固定内存中分配大量张量或较大的张量可能会占用大量RAM。 这会减少其他关键操作（例如分页）可用的内存，从而对算法的整体性能产生负面影响。

结论

在整个教程中，我们探讨了几个关键因素，这些因素影响着从主机发送张量到设备时的速度和内存管理。我们了解到，使用 non_blocking=True 通常可以加速数据传输，而 pin_memory() 如果正确实现也可以提升性能。然而，这些技术需要精心设计和校准才能有效。

记住，优化代码资源使用和实现最佳性能的关键在于对代码进行 profiling，并关注内存消耗。

附加资源

如果您在使用CUDA设备时遇到内存复制问题，或者想了解更多本教程中讨论的内容，请参阅以下参考文献：

• CUDA 工具包内存管理文档;

• CUDA 固定内存注意事项;

• 如何在CUDA C/C++中优化数据传输;

• tensordict 文档 和 仓库。