了解 CUDA 内存使用情况¶

为了调试 CUDA 内存使用情况，PyTorch 提供了一种方法来生成内存快照，这些快照记录了任何时间点已分配的 CUDA 内存的状态，并可以选择记录导致该快照的分配事件的历史记录。

然后，可以将生成的快照拖放到托管在 pytorch.org/memory_viz 上的交互式查看器中，可以使用该查看器来浏览快照。

生成快照¶

记录快照的常见模式是启用内存历史记录，运行要观察的代码，然后保存包含已腌制快照的文件。

# enable memory history, which will
# add tracebacks and event history to snapshots
torch.cuda.memory._record_memory_history()

run_your_code()
torch.cuda.memory._dump_snapshot("my_snapshot.pickle")

使用可视化工具¶

打开 pytorch.org/memory_viz 并将腌制快照文件拖放到可视化工具中。可视化工具是一个在您的计算机上本地运行的 JavaScript 应用程序。它不会上传任何快照数据。

活动内存时间线¶

活动内存时间线显示特定 GPU 上快照中所有实时张量随时间推移的变化。在绘图上平移/缩放以查看较小的分配。将鼠标悬停在已分配的块上以查看分配该块时的堆栈跟踪以及地址等详细信息。可以调整详细信息滑块以渲染较少的分配，并在数据量很大时提高性能。

分配器状态历史记录¶

分配器状态历史记录在左侧时间线上显示单个分配器事件。选择时间线中的一个事件以查看该事件时分配器状态的视觉摘要。该摘要显示从 cudaMalloc 返回的每个单独段以及如何将其分成单个分配的块或空闲空间。将鼠标悬停在段和块上以查看分配内存时的堆栈跟踪。将鼠标悬停在事件上以查看事件发生时的堆栈跟踪，例如当张量被释放时。内存不足错误报告为 OOM 事件。查看 OOM 期间的内存状态可能会提供有关为什么即使保留内存仍然存在分配失败的见解。

堆栈跟踪信息还报告分配发生的地址。地址 b7f064c000000_0 指的是地址为 7f064c000000 的（b）块，这是此地址第“_0”次分配。可以在活动内存时间线中查找此唯一字符串，并在活动状态历史记录中搜索它以检查分配或释放张量时的内存状态。

快照 API 参考¶

torch.cuda.memory._record_memory_history(enabled='all', context='all', stacks='all', max_entries=9223372036854775807, device=None)[source]¶

启用记录与内存分配相关的堆栈跟踪，以便您可以了解在 torch.cuda.memory._snapshot() 中分配了哪段内存。

除了保存当前每个分配和释放的堆栈跟踪之外，这还将启用对所有分配/释放事件的历史记录的记录。

使用 torch.cuda.memory._snapshot() 检索此信息，并使用 _memory_viz.py 中的工具可视化快照。

Python 跟踪收集速度很快（每次跟踪 2 微秒），因此如果您预期将来可能需要调试内存问题，可以考虑在生产作业中启用它。

C++ 跟踪收集速度也很快（每帧约 50 纳秒），对于许多典型的程序，这相当于每次跟踪约 2 微秒，但会根据堆栈深度而有所不同。

参数

enabled (Literal[None, "state", "all"], optional) – None，禁用记录内存历史记录。 “state”，保留当前已分配内存的信息。 “all”，此外还保留所有分配/释放调用的历史记录。默认值为“all”。
context (Literal[None, "state", "alloc", "all"], 可选) – None，不记录任何跟踪信息。 “state”，记录当前分配内存的跟踪信息。 “alloc”，此外还保留分配调用的跟踪信息。 “all”，此外还保留释放调用的跟踪信息。默认值为“all”。
stacks (Literal["python", "all"], 可选) – “python”，在跟踪信息中包含 Python、TorchScript 和 inductor 帧 “all”，此外还包含 C++ 帧，默认值为“all”。
max_entries (int, 可选) – 在记录的历史记录中，最多保留 max_entries 个分配/释放事件。

torch.cuda.memory._snapshot(device=None)[source]¶

保存调用该函数时 CUDA 内存状态的快照。

该状态表示为一个字典，结构如下。

class Snapshot(TypedDict):
    segments : List[Segment]
    device_traces: List[List[TraceEntry]]

class Segment(TypedDict):
    # Segments are memory returned from a cudaMalloc call.
    # The size of reserved memory is the sum of all Segments.
    # Segments are cached and reused for future allocations.
    # If the reuse is smaller than the segment, the segment
    # is split into more then one Block.
    # empty_cache() frees Segments that are entirely inactive.
    address: int
    total_size: int #  cudaMalloc'd size of segment
    stream: int
    segment_type: Literal['small', 'large'] # 'large' (>1MB)
    allocated_size: int # size of memory in use
    active_size: int # size of memory in use or in active_awaiting_free state
    blocks : List[Block]

class Block(TypedDict):
    # A piece of memory returned from the allocator, or
    # current cached but inactive.
    size: int
    requested_size: int # size requested during malloc, may be smaller than
                        # size due to rounding
    address: int
    state: Literal['active_allocated', # used by a tensor
                'active_awaiting_free', # waiting for another stream to finish using
                                        # this, then it will become free
                'inactive',] # free for reuse
    frames: List[Frame] # stack trace from where the allocation occurred

class Frame(TypedDict):
        filename: str
        line: int
        name: str

class TraceEntry(TypedDict):
    # When `torch.cuda.memory._record_memory_history()` is enabled,
    # the snapshot will contain TraceEntry objects that record each
    # action the allocator took.
    action: Literal[
    'alloc'  # memory allocated
    'free_requested', # the allocated received a call to free memory
    'free_completed', # the memory that was requested to be freed is now
                    # able to be used in future allocation calls
    'segment_alloc', # the caching allocator ask cudaMalloc for more memory
                    # and added it as a segment in its cache
    'segment_free',  # the caching allocator called cudaFree to return memory
                    # to cuda possibly trying free up memory to
                    # allocate more segments or because empty_caches was called
    'oom',          # the allocator threw an OOM exception. 'size' is
                    # the requested number of bytes that did not succeed
    'snapshot'      # the allocator generated a memory snapshot
                    # useful to coorelate a previously taken
                    # snapshot with this trace
    ]
    addr: int # not present for OOM
    frames: List[Frame]
    size: int
    stream: int
    device_free: int # only present for OOM, the amount of
                    # memory cuda still reports to be free

返回值: 快照字典对象

torch.cuda.memory._dump_snapshot(filename='dump_snapshot.pickle')[source]¶

将 torch.memory._snapshot() 字典的腌制版本保存到文件。

此文件可以通过 pytorch.org/memory_viz 上的交互式快照查看器打开。

参数: filename (str, 可选) – 要创建的文件的名称。默认值为“dump_snapshot.pickle”。