get_memory_stats

torchtune.training.get_memory_stats(device: device, reset_stats: bool = True) → dict

计算传递给该设备的内存摘要。如果 reset_stats 为 True，这也将重置 CUDA 的峰值内存跟踪。这对于获取有关峰值内存相对使用情况的数据（例如，模型初始化期间、正向传播期间等的峰值内存）以及优化内存以用于训练的各个部分非常有用。

参数：

• device (torch.device) – 获取内存摘要的设备。仅支持 CUDA 设备。

• reset_stats (bool) – 是否重置 CUDA 的峰值内存跟踪。

返回值：

包含峰值活动内存、峰值分配内存和峰值保留内存的字典。此字典对于记录内存统计信息很有用。

返回类型：

Dict[str, float]

引发：

ValueError – 如果传入的设备不是 CUDA。