get_memory_usage() — pytorch Function Reference

Architecture documentation for the get_memory_usage() function in dataloader_benchmark.py from the pytorch codebase.

Function python CoreTensor Dispatcher called by 1

Entity Profile

CoreTensor→ Dispatcher→ get_memory_usage() — pytorch Function Reference

Dependency Diagram

graph TD
  8002ad3e_d000_9ebf_bd29_6dd9e7ed2ade["get_memory_usage()"]
  f6150286_ce63_72f2_12ae_6ae1460174f5["benchmark_dataloader()"]
  f6150286_ce63_72f2_12ae_6ae1460174f5 -->|calls| 8002ad3e_d000_9ebf_bd29_6dd9e7ed2ade
  style 8002ad3e_d000_9ebf_bd29_6dd9e7ed2ade fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/data/dataloader_benchmark.py lines 31–53

def get_memory_usage():
    """
    Get current memory usage in MB. This includes all child processes.

    Returns:
        Total memory usage in MB
    """
    process = psutil.Process()

    main_memory = process.memory_full_info().pss

    # Add memory usage of all child processes
    for child in process.children(recursive=True):
        try:
            child_mem = child.memory_full_info().pss
            main_memory += child_mem
        except (psutil.NoSuchProcess, psutil.AccessDenied, AttributeError):
            # Process might have terminated or doesn't support PSS, fall back to USS
            print(f"Failed to get PSS for {child}, falling back to USS")
            child_mem = child.memory_info().uss
            main_memory += child_mem

    return main_memory / (1024 * 1024)