理解 CUDA 記憶體使用¶

為了除錯 CUDA 記憶體使用情況，PyTorch 提供了一種生成記憶體快照的方法，可以記錄任意時間點已分配 CUDA 記憶體的狀態，並可選擇記錄導致該快照的分配事件歷史。

生成的資料可以拖放到 pytorch.org/memory_viz 上託管的互動式檢視工具中，用於探索快照。

生成快照¶

記錄快照的常見模式是：啟用記憶體歷史記錄，執行待觀察的程式碼，然後將序列化（pickled）的快照儲存到檔案中。

# 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 並將序列化（pickled）的快照檔案拖放到視覺化工具中。該視覺化工具是一個在您計算機本地執行的 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][source]¶

啟用記錄與記憶體分配相關的堆疊跟蹤，以便您知道 torch.cuda.memory._snapshot() 中任何一塊記憶體是由什麼分配的。

除了保留每次當前分配和釋放的堆疊跟蹤外，此功能還將啟用記錄所有分配/釋放事件的歷史記錄。

使用 torch.cuda.memory._snapshot() 檢索此資訊，並使用 _memory_viz.py 中的工具來視覺化快照。

Python 跟蹤收集速度很快（每次跟蹤 2 微秒），因此如果您預計需要除錯記憶體問題，可以考慮在生產作業中啟用此功能。

C++ 跟蹤收集速度也很快（約 50 納秒/幀），對於許多典型程式來說，這相當於每次跟蹤約 2 微秒，但可能因堆疊深度而異。

引數

enabled (Literal[None, "state", "all"], 可選) – 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][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][source]¶

將 torch.memory._snapshot() 字典的序列化（pickled）版本儲存到檔案中。

該檔案可以使用 pytorch.org/memory_viz 上的互動式快照檢視工具開啟。

引數: filename (str, 可選) – 要建立的檔名。預設為 “dump_snapshot.pickle”。