查看PyTorch的GPU使用情况的工具

文章目录

torch.cuda API
Snapshot
torchinfo

torch.cuda API

torch.cuda.memory_stats：返回给定设备的 CUDA 内存分配器统计信息字典。该函数的返回值是一个统计字典，每个字典都是一个非负整数。
torch.cuda.memory_summary：返回给定设备当前内存分配器统计信息的人类可读的打印输出。这对于在训练期间或处理内存不足异常时定期显示非常有用。
torch.cuda.memory_snapshot：返回所有设备上的 CUDA 内存分配器状态的快照。解释该函数的输出需要熟悉内存分配器的内部结构。
torch.cuda.reset_peak_memory_stats：重置 CUDA 内存分配器跟踪的“峰值”统计数据。峰值统计数据对应于每个单独统计字典中的“峰值”键。
torch.cuda.memory_allocated：返回给定设备的张量当前占用的 GPU 内存（以字节为单位）。这可能小于 nvidia-smi 中显示的数量，因为缓存分配器可以保留一些未使用的内存，并且需要在 GPU 上创建一些上下文。
torch.cuda.max_memory_allocated：返回给定设备的张量占用的最大 GPU 内存（以字节为单位）。默认情况下，这将返回自该程序开始以来分配的内存峰值。reset_peak_memory_stats()可用于重置跟踪此指标的起点。例如，这两个函数可以测量训练循环中每次迭代的峰值分配内存使用情况。
torch.cuda.memory_reserved：返回给定设备的缓存分配器管理的当前 GPU 内存（以字节为单位）。
torch.cuda.max_memory_reserved：返回给定设备的缓存分配器管理的最大 GPU 内存（以字节为单位）。默认情况下，这将返回自该程序开始以来的峰值缓存内存。reset_peak_memory_stats() 可用于重置跟踪此指标的起点。例如，这两个函数可以测量训练循环中每次迭代的峰值缓存内存量。
torch.cuda.empty_cache：释放缓存分配器当前持有的所有未占用的缓存内存，以便这些内存可以在其他 GPU 应用程序中使用并在 nvidia-smi 中可见。empty_cache()不会增加 PyTorch 可用的 GPU 内存量。但是，在某些情况下，它可能有助于减少 GPU 内存碎片。

以下内容翻译自 Memory management 文档

PyTorch使用缓存内存分配器 (caching memory allocator) 来加快内存分配速度。这允许快速释放内存而无需设备同步。然而，由分配器管理的未使用内存在nvidia-smi中仍会显示为已使用。

你可以使用memory_allocated()和max_memory_allocated()来监控由张量占用的内存，
并使用memory_reserved()和max_memory_reserved()来监控缓存分配器管理的总内存量。
调用empty_cache()可以释放PyTorch中所有未使用的缓存内存，以便其他GPU应用程序可以使用这些内存。然而，由张量占用的GPU内存不会被释放，因此它不能增加PyTorch可用的GPU内存量。

对于更高级的用户，我们提供了通过memory_stats()进行更全面的内存基准测试的功能。我们还提供了通过memory_snapshot()捕获内存分配器状态的完整快照的功能，这可以帮助你理解代码产生的底层分配模式。

在环境中设置PYTORCH_NO_CUDA_MEMORY_CACHING=1以禁用缓存。

缓存分配器的行为可以通过环境变量PYTORCH_CUDA_ALLOC_CONF来控制。格式为PYTORCH_CUDA_ALLOC_CONF=<option>:<value>,<option2>:<value2>...可用的选项有：

max_split_size_mb 阻止分配器拆分大于此大小（以MB为单位）的块。这可以帮助防止碎片化，并可能允许一些边缘工作负载在不耗尽内存的情况下完成。性能成本可能从‘零’到‘显著’不等，具体取决于分配模式。默认值是无限制，即所有块都可以被拆分。此选项应作为工作负载因‘内存不足’而中止并显示大量非活动拆分块时的最后手段。
roundup_power2_divisions 帮助将请求的分配大小四舍五入到最接近的2的幂除法，以更好地利用块。在当前的CUDACachingAllocator中，大小被四舍五入为512的块大小的倍数，因此这对于较小的大小来说很好。然而，对于较大的近邻分配，这可能效率低下，因为每个分配都会去不同大小的块，减少了这些块的重用。这可能会创建大量未使用的块，浪费GPU内存容量。此选项使分配大小四舍五入到最接近的2的幂除法。例如，如果我们需要四舍五入大小为1200，并且分割数为4，1200大小介于1024和2048之间，如果我们在它们之间进行4次分割，值分别为1024、1280、1536和1792。因此，分配大小为1200将被四舍五入到1280作为最接近的2的幂分割上限。
roundup_bypass_threshold_mb 对于超过阈值大小（以MB为单位）的分配请求，绕过四舍五入分配大小。这可以在进行预计会持久或寿命较长的大型分配时帮助减少内存占用。
garbage_collection_threshold 帮助主动回收未使用的GPU内存，以避免触发代价高昂的同步和回收所有操作（release_cached_blocks），这对延迟敏感的GPU应用程序（如服务器）可能不利。设置此阈值后（例如0.8），如果GPU内存容量使用超过阈值（即分配给GPU应用程序的总内存的80%），分配器将开始回收GPU内存块。算法倾向于首先释放旧的和未使用的块，以避免释放正在积极重用的块。阈值值应大于0.0且小于1.0。

Snapshot

以下内容摘自博客：PyTorch显存可视化与Snapshot数据分析

以下内容摘自博客：Pytorch 显存管理机制与显存占用分析方法

torchinfo

torchinfo，可实现模型参数量计算、各层特征图形状计算和计算量计算等功能。

使用参考：6.5 模型参数打印 · PyTorch实用教程(第二版)

支持版本：PyTorch 1.4.0+

安装：pip install torchinfo

def summary(
    model: nn.Module,
    input_size: Optional[INPUT_SIZE_TYPE] = None,
    input_data: Optional[INPUT_DATA_TYPE] = None,
    batch_dim: Optional[int] = None,
    cache_forward_pass: Optional[bool] = None,
    col_names: Optional[Iterable[str]] = None,
    col_width: int = 25,
    depth: int = 3,
    device: Optional[torch.device] = None,
    dtypes: Optional[List[torch.dtype]] = None,
    mode: str | None = None,
    row_settings: Optional[Iterable[str]] = None,
    verbose: int = 1,
    **kwargs: Any,
) -> ModelStatistics:

Summarize the given PyTorch model. Summarized information includes: 总结给定的 PyTorch 模型。汇总信息包括：

Layer names,
input/output shapes,
kernel shape,
# of parameters,
# of operations (Mult-Adds),
whether layer is trainable

可选参数：

model (nn.Module)：要汇总的 PyTorch 模型。模型应该完全处于 train() 或 eval() 模式。如果层不全处于同一模式，运行 summary 可能会对 batchnorm 或 dropout 统计数据产生副作用。
input_size (Sequence of Sizes)：输入数据的形状，格式为 List/Tuple/torch.Size（数据类型必须与模型输入匹配，默认是 FloatTensors）。应在元组中包括批量大小。默认值：None
input_data (Sequence of Tensors)：模型前向传递的参数（数据类型推断）。如果 forward() 函数需要多个参数，请传入 args 列表或 kwargs 字典（如果 forward() 函数只接受一个字典参数，请将其包装在一个列表中）。默认值：None
batch_dim (int)：输入数据的批量维度。如果 batch_dim 为 None，则假设 input_data / input_size 包含批量维度，并在所有计算中使用它。否则，扩展所有张量以包含 batch_dim。指定 batch_dim 可以作为运行时优化，因为如果指定了 batch_dim，torchinfo 在前向传递时使用批量大小为 1。默认值：None
cache_forward_pass (bool)：如果为 True，使用模型类名称作为键缓存 forward() 函数的运行。如果前向传递是一个耗时的操作，这可以更容易地修改模型摘要的格式，例如更改深度或启用的列类型，特别是在 Jupyter Notebooks 中。警告：在启用此功能时修改模型架构或输入数据/输入大小不会使缓存失效或重新运行前向传递，可能会导致不正确的摘要。默认值：False
col_names (Iterable[str])：指定输出中要显示的列。目前支持的有：“input_size”, “output_size”, “num_params”, “params_percent”, “kernel_size”, “mult_adds”, “trainable”。默认值：“output_size”, “num_params”。如果未提供 input_data / input_size，则仅使用 “num_params”。
col_width (int)：每列的宽度。默认值：25
depth (int)：显示嵌套层的深度（例如，Sequentials）。低于此深度的嵌套层将不会在摘要中显示。默认值：3
device (torch.Device)：为模型和 input_data 使用此 torch 设备。如果未指定，则使用 input_data 的数据类型（如果已提供）或模型的参数。否则，使用 torch.cuda.is_available() 的结果。默认值：None
dtypes (List[torch.dtype])：如果使用 input_size，torchinfo 假设输入使用 FloatTensors。如果模型使用不同的数据类型，请指定该数据类型。对于多个输入，请指定每个输入的大小，并在此处指定每个参数的类型。默认值：None
mode (str)：“train” 或 “eval” 中的一个，决定在调用 summary() 之前调用 model.train() 还是 model.eval()。默认值：“eval”
row_settings (Iterable[str])：指定在一行中显示哪些特性。目前支持的有：“ascii_only”, “depth”, “var_names”。默认值：“depth”
verbose (int)：0（静默）：无输出；1（默认）：打印模型摘要；2（详细）：详细显示权重和偏置层。默认值：1。如果使用 Jupyter Notebook 或 Google Colab，默认值为 0。
**kwargs：模型 forward 函数中使用的其他参数。不再支持传递 *args。