torch.profiler

概觀

PyTorch Profiler 是一個允許在訓練和推論期間收集效能指標的工具。 Profiler 的上下文管理器 API 可用於更好地瞭解哪些模型運算子的成本最高，檢查其輸入形狀和堆疊追蹤，研究裝置核心活動並視覺化執行追蹤。

注意

torch.autograd 模組中較早版本的 API 被視為舊版，將被棄用。

API 參考

class torch.profiler._KinetoProfile(*, activities=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, with_modules=False, experimental_config=None, execution_trace_observer=None)

低階分析器包裝 autograd 分析器

參數

• activities (iterable) – 要在分析中使用的活動群組 (CPU、CUDA) 清單，支援的值：torch.profiler.ProfilerActivity.CPU、torch.profiler.ProfilerActivity.CUDA、torch.profiler.ProfilerActivity.XPU。預設值：ProfilerActivity.CPU 和（可用時）ProfilerActivity.CUDA 或（可用時）ProfilerActivity.XPU。

• record_shapes (bool) – 儲存有關運算子輸入形狀的資訊。

• profile_memory (bool) – 追蹤張量記憶體配置/解除配置（有關詳細資訊，請參閱 export_memory_timeline）。

• with_stack (bool) – 記錄作業的原始程式碼資訊（檔案和行號）。

• with_flops (bool) – 使用公式來估計特定運算子（矩陣乘法和 2D 卷積）的 FLOPS。

• with_modules (bool) – 記錄與作業的呼叫堆疊相對應的模組階層結構（包括函數名稱）。例如，如果模組 A 的正向呼叫模組 B 的正向呼叫包含 aten::add 作業，則 aten::add 的模組階層結構為 A.B。請注意，目前僅 TorchScript 模型支援此功能，而 Eager 模式模型則不支援。

• experimental_config (_ExperimentalConfig) – 由 Kineto 等分析器函式庫使用的一組實驗性選項。請注意，不保證向後相容性。

• execution_trace_observer (ExecutionTraceObserver) – PyTorch 執行追蹤觀察器物件。 PyTorch 執行追蹤 提供 AI/ML 工作負載的基於圖形的表示，並啟用重播基準測試、模擬器和模擬器。當包含此參數時，觀察器的 start() 和 stop() 將在與 PyTorch 分析器相同的時間範圍內被呼叫。

注意

此 API 處於實驗階段，未來可能會有所變更。

啟用形狀和堆疊追蹤會導致額外的開銷。當指定 record_shapes=True 時，分析器將暫時保留對張量的引用；這可能會進一步阻止依賴引用計數的某些最佳化，並引入額外的張量副本。

add_metadata(key, value)

使用字串鍵和字串值將使用者定義的中繼資料新增至追蹤檔案

add_metadata_json(key, value)

使用字串鍵和有效的 json 值將使用者定義的中繼資料新增至追蹤檔案

events()

傳回未彙總的分析器事件清單，可在追蹤回呼中或分析完成後使用

export_chrome_trace(path)

以 Chrome JSON 格式匯出收集的追蹤。如果已啟用 kineto，則只會匯出排程中的最後一個週期。

export_memory_timeline(path, device=None)

匯出指定裝置的分析器收集樹中的記憶體事件資訊，並匯出時間軸圖表。使用 export_memory_timeline 可以匯出 3 個檔案，每個檔案都由 path 的後綴控制。

• 對於相容 HTML 的圖表，請使用後綴 .html，記憶體時間軸圖表將作為 PNG 檔案嵌入 HTML 檔案中。

• 對於由 [times, [sizes by category]] 組成的圖表點，其中 times 是時間戳記，而 sizes 是每個類別的記憶體使用量。記憶體時間軸圖表將根據後綴儲存為 JSON (.json) 或 gzip 壓縮的 JSON (.json.gz)。

• 對於原始記憶體點，請使用後綴 .raw.json.gz。每個原始記憶體事件將包含 (timestamp, action, numbytes, category)，其中 action 是 [PREEXISTING, CREATE, INCREMENT_VERSION, DESTROY] 其中之一，而 category 是 torch.profiler._memory_profiler.Category 中的其中一個列舉。

輸出：記憶體時間軸以 gzip 壓縮的 JSON、JSON 或 HTML 格式寫入。

export_stacks(path, metric='self_cpu_time_total')

將堆疊追蹤儲存到檔案

參數

• path (str) – 將堆疊檔案儲存到此位置；

• metric (str) – 要使用的指標：「self_cpu_time_total」或「self_cuda_time_total」

key_averages(group_by_input_shape=False, group_by_stack_n=0)

計算事件的平均值，並根據運算子名稱和（可選）輸入形狀和堆疊對其進行分組。

注意

若要使用形狀/堆疊功能，請確保在建立分析器環境管理員時設定 record_shapes/with_stack。

preset_metadata_json(key, value)

在分析器尚未啟動且稍後新增到追蹤檔案中時，預設使用者定義的中繼資料。中繼資料的格式為字串金鑰和有效的 json 值

class torch.profiler.profile(*, activities=None, schedule=None, on_trace_ready=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, with_modules=False, experimental_config=None, execution_trace_observer=None, use_cuda=None)

分析器環境管理員。

參數

• activities (iterable) – 要在分析中使用的活動群組 (CPU、CUDA) 清單，支援的值：torch.profiler.ProfilerActivity.CPU、torch.profiler.ProfilerActivity.CUDA、torch.profiler.ProfilerActivity.XPU。預設值：ProfilerActivity.CPU 和（可用時）ProfilerActivity.CUDA 或（可用時）ProfilerActivity.XPU。

• schedule (Callable) – 將步驟 (int) 作為單一參數並傳回 ProfilerAction 值的可呼叫函式，該值指定要在每個步驟執行的分析器動作。

• on_trace_ready (Callable) – 在分析期間，當 schedule 傳回 ProfilerAction.RECORD_AND_SAVE 時，會在每個步驟呼叫的可呼叫函式。

• record_shapes (bool) – 儲存有關運算子輸入形狀的資訊。

• profile_memory (bool) – 追蹤張量記憶體配置/解除配置。

• with_stack (bool) – 記錄作業的原始程式碼資訊（檔案和行號）。

• with_flops (bool) – 使用公式來估計特定運算子（矩陣乘法和 2D 卷積）的 FLOPs（浮點運算）。

• with_modules (bool) – 記錄與作業的呼叫堆疊相對應的模組階層結構（包括函數名稱）。例如，如果模組 A 的正向呼叫模組 B 的正向呼叫包含 aten::add 作業，則 aten::add 的模組階層結構為 A.B。請注意，目前僅 TorchScript 模型支援此功能，而 Eager 模式模型則不支援。

• experimental_config (_ExperimentalConfig) – 一組用於 Kineto 程式庫功能的實驗性選項。請注意，不保證向後相容性。

• execution_trace_observer (ExecutionTraceObserver) – PyTorch 執行追蹤觀察器物件。 PyTorch 執行追蹤 提供了 AI/ML 工作負載的基於圖表的表示，並啟用了重播基準測試、模擬器和模擬器。當包含此參數時，將針對與 PyTorch 分析器相同的時間範圍呼叫觀察器的 start() 和 stop()。如需程式碼範例，請參閱下面的範例部分。

• use_cuda (bool) –

  自版本 1.8.1 起已棄用： 請改用 activities。

注意

使用 schedule() 來產生可呼叫的排程。在分析長時間訓練作業時，非預設排程很有用，並且允許使用者在訓練過程的不同迭代中取得多個追蹤。預設排程只是在環境管理員的持續時間內持續記錄所有事件。

注意

使用 tensorboard_trace_handler() 為 TensorBoard 產生結果檔案

on_trace_ready=torch.profiler.tensorboard_trace_handler(dir_name)

分析後，可以在指定的目錄中找到結果檔案。使用以下命令

tensorboard --logdir dir_name

在 TensorBoard 中查看結果。如需詳細資訊，請參閱 PyTorch Profiler TensorBoard 外掛

注意

啟用形狀和堆疊追蹤會導致額外的開銷。當指定 record_shapes=True 時，分析器將暫時保留對張量的引用；這可能會進一步阻止依賴引用計數的某些最佳化，並引入額外的張量副本。

範例

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ]
) as p:
    code_to_profile()
print(p.key_averages().table(
    sort_by="self_cuda_time_total", row_limit=-1))

使用分析器的 schedule、on_trace_ready 和 step 函式

# Non-default profiler schedule allows user to turn profiler on and off
# on different iterations of the training loop;
# trace_handler is called every time a new trace becomes available
def trace_handler(prof):
    print(prof.key_averages().table(
        sort_by="self_cuda_time_total", row_limit=-1))
    # prof.export_chrome_trace("/tmp/test_trace_" + str(prof.step_num) + ".json")

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],

    # In this example with wait=1, warmup=1, active=2, repeat=1,
    # profiler will skip the first step/iteration,
    # start warming up on the second, record
    # the third and the forth iterations,
    # after which the trace will become available
    # and on_trace_ready (when set) is called;
    # the cycle repeats starting with the next step

    schedule=torch.profiler.schedule(
        wait=1,
        warmup=1,
        active=2,
        repeat=1),
    on_trace_ready=trace_handler
    # on_trace_ready=torch.profiler.tensorboard_trace_handler('./log')
    # used when outputting for tensorboard
    ) as p:
        for iter in range(N):
            code_iteration_to_profile(iter)
            # send a signal to the profiler that the next iteration has started
            p.step()

以下範例顯示如何設定執行追蹤觀察器 (execution_trace_observer)

with torch.profiler.profile(
    ...
    execution_trace_observer=(
        ExecutionTraceObserver().register_callback("./execution_trace.json")
    ),
) as p:
    for iter in range(N):
        code_iteration_to_profile(iter)
        p.step()

您也可以參閱 tests/profiler/test_profiler.py 中的 test_execution_trace_with_kineto()。注意：您也可以傳遞任何滿足 _ITraceObserver 介面的物件。

step()

向分析器發出訊號，表示下一個分析步驟已開始。

class torch.profiler.ProfilerAction(value)

可以在指定間隔採取的分析器動作

class torch.profiler.ProfilerActivity

成員

CPU

XPU

MTIA

CUDA

PrivateUse1

property name

torch.profiler.schedule(*, wait, warmup, active, repeat=0, skip_first=0)

回傳一個可呼叫的物件，可作為分析器 schedule 參數。分析器將會略過前 skip_first 個步驟，然後等待 wait 個步驟，接著對接下來的 warmup 個步驟進行預熱，然後對接下來的 active 個步驟進行主動記錄，最後從 wait 個步驟開始重複此循環。可選的循環次數由 repeat 參數指定，零值表示循環將持續到分析完成。

回傳型別

可呼叫物件

torch.profiler.tensorboard_trace_handler(dir_name, worker_name=None, use_gzip=False)

將追蹤檔案輸出到 dir_name 目錄，然後該目錄可以直接作為日誌目錄傳遞給 TensorBoard。 worker_name 在分散式情境中，每個工作者的名稱都應該不同，預設會設定為「[主機名稱]_[進程識別碼]」。

英特爾儀器與追蹤技術 API

torch.profiler.itt.is_available()

檢查 ITT 功能是否可用

torch.profiler.itt.mark(msg)

描述在某個時間點發生的瞬間事件。

參數

msg (str) – 與事件關聯的 ASCII 訊息。

torch.profiler.itt.range_push(msg)

將一個範圍推入巢狀範圍跨度的堆疊。回傳開始範圍的從零開始的深度。

參數

msg (str) – 與範圍關聯的 ASCII 訊息

torch.profiler.itt.range_pop()

從巢狀範圍跨度的堆疊中彈出一個範圍。回傳結束範圍的從零開始的深度。