【Python】VRAM使用量を調べるデコレーターをClaude 3.5 Sonnetに書いてもらいました。

注意

動作は問題なくします。

Windowsタスクマネージャーが表示するVRAM使用量と一致していません。

どちらが正確かはわかりません。

Pythonスクリプト

# pip install pynvml

import functools
import threading
import time
from pynvml import *
import torch
import gc

def gpu_monitor(interval=1):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            stop_event = threading.Event()

            def monitor_thread():
                max_vram = 0
                device_id = 0
                nvmlInit()
                handle = nvmlDeviceGetHandleByIndex(device_id)
                info = nvmlDeviceGetMemoryInfo(handle)
                total_vram = info.total / (1024**3)
                try:
                    while not stop_event.is_set():
                        info = nvmlDeviceGetMemoryInfo(handle)
                        using_vram = info.used / (1024**3)
                        if using_vram > max_vram:
                            max_vram = using_vram
                        
                        time.sleep(interval)
                except NVMLError as error:
                    print(f"NVML Error: {error}")
                finally:
                    nvmlShutdown()
                    print(f"GPU {device_id} - Used memory: {max_vram:.2f}/{total_vram:.2f} GB")

            monitor = threading.Thread(target=monitor_thread)
            monitor.start()

            try:
                result = func(*args, **kwargs)
            finally:
                stop_event.set()
                monitor.join()

            return result
        return wrapper
    return decorator

def time_monitor(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        print(f"time: {(end_time - start_time):.2f} sec")
        return result
    return wrapper

def reset_memory():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_accumulated_memory_stats()
    torch.cuda.reset_peak_memory_stats()

def print_memory(mode: int = 0):
    if mode == 1:
        memory = round(torch.cuda.memory_allocated() / 1024**3, 2)
        print(f"{memory=} GB")
    max_memory = round(torch.cuda.max_memory_allocated() / 1024**3, 2)
    max_reserved = round(torch.cuda.max_memory_reserved() / 1024**3, 2)
    print(f"{max_memory=} GB")
    print(f"{max_reserved=} GB")

使い方の例

import time
import torch
from diffusers import FluxTransformer2DModel, FluxPipeline
from transformers import T5EncoderModel

from monitor_vram import gpu_monitor

@gpu_monitor(interval=0.5)
def main():
    start = time.perf_counter()

    transformer = FluxTransformer2DModel.from_pretrained(
        "FLUX.1-transformer-nf4",
        torch_dtype=torch.bfloat16
    )
    text_encoder_2 = T5EncoderModel.from_pretrained(
        "FLUX.1-text_encoder_2-nf4",
        torch_dtype=torch.bfloat16
    )

    pipe = FluxPipeline.from_pretrained(
        "FLUX.1-dev",
        transformer=transformer,
        text_encoder_2=text_encoder_2,
        torch_dtype=torch.bfloat16
    )

    pipe.enable_model_cpu_offload()

    generator = torch.Generator().manual_seed(123)
    prompt = "A photorealistic portrait of a young Japanese woman with long black hair and natural makeup, wearing a casual white blouse, sitting in a modern Tokyo cafe with soft window light"
    out = pipe(
        prompt=prompt,
        height=768,
        width=1360,
        num_inference_steps=50,
        generator=generator,
        guidance_scale=3.5
    ).images[0]
    out.save("girl.jpg")

    end = time.perf_counter()
    print(f"time: {(end - start):.2f}sec")

if __name__ == "__main__":
    main()

Windowsタスクマネージャーとの違い

これもClaude 3.5 Sonnetに聞いてみました。

これは興味深い質問ですね。pynvmlライブラリのnvmlDeviceGetMemoryInfoとWindowsのタスクマネージャーで表示されるVRAM使用量に違いがある理由はいくつか考えられます:

測定方法の違い:
pynvmlは、NVIDIA GPUのドライバーから直接情報を取得します。
タスクマネージャーは、Windows OSを通じて情報を取得します。

測定するメモリの種類:
pynvmlは、通常、GPUの「専用メモリ」使用量を報告します。
タスクマネージャーは、「共有メモリ」を含む場合があります。

キャッシュやバッファの扱い:
pynvmlは、キャッシュやバッファメモリを使用中のメモリとして計上する可能性があります。
タスクマネージャーは、これらを別途扱う場合があります。

更新頻度:
pynvmlは、リアルタイムに近い情報を提供します。
タスクマネージャーは、更新頻度が低い場合があります。

ドライバーやソフトウェアの違い:
NVIDIAドライバーのバージョンや、使用しているソフトウェアによって、報告される使用量が異なる場合があります。

これらの違いにより、二つの方法で取得したVRAM使用量に差が生じる可能性があります。より正確なGPUメモリ使用量を知りたい場合は、通常pynvmlやNVIDIA-SMIなどのNVIDIA公式ツールの方が信頼性が高いとされています。