算力探查器
DeepSpeed 中的算力探查器会分析模型的前向传播,并测量其参数、延迟和浮点运算。DeepSpeed 算力探查器可以与 DeepSpeed 运行时一起使用,也可以作为独立的软件包使用。
当使用 DeepSpeed 进行模型训练时,可以在 deepspeed_config 文件中配置算力探查器,而无需更改用户代码。要将算力探查器用于 DeepSpeed 运行时之外,只需安装 DeepSpeed 并导入 flops_profiler 软件包即可直接使用 API。
有关使用详细信息,请参阅 算力探查器教程。
算力探查器
- class deepspeed.profiling.flops_profiler.profiler.FlopsProfiler(model, ds_engine=None, recompute_fwd_factor=0.0)[source]
基类:
object
测量 PyTorch 模型中每个模块的延迟、估计的浮点运算次数和参数。
算力探查器会分析 PyTorch 模型的前向传播,并打印模型图,其中每个模块都附加了测量的分析结果。它显示了模型中如何花费延迟、算力和参数,以及哪些模块或层可能是瓶颈。它还会输出根据聚合延迟、算力和参数在深度 l 处的前 k 个模块的名称,其中 k 和 l 由用户指定。对于每一批输入,都会计算输出分析结果。DeepSpeed 算力探查器可以与 DeepSpeed 运行时一起使用,也可以作为独立的软件包使用。当使用 DeepSpeed 进行模型训练时,可以在 deepspeed_config 文件中配置算力探查器,无需更改用户代码。
如果将探查器用作独立软件包,则导入 flops_profiler 软件包并使用 API。
以下是在典型训练工作流程中使用的一个示例
model = Model() prof = FlopsProfiler(model) for step, batch in enumerate(data_loader): if step == profile_step: prof.start_profile() loss = model(batch) if step == profile_step: flops = prof.get_total_flops(as_string=True) params = prof.get_total_params(as_string=True) prof.print_model_profile(profile_step=profile_step) prof.end_profile() loss.backward() optimizer.step()
要分析推理中训练好的模型,请使用 get_model_profile API。
- 参数
object (torch.nn.Module) – 要分析的 PyTorch 模型。
- start_profile(ignore_list=None)[source]
开始分析。
会递归地向所有模块添加额外的属性,并且会修补分析过的 torch.nn.functionals。
- 参数
ignore_list (list, 可选) – 分析时要忽略的模块列表。默认为 None。
- get_total_flops(as_string=False)[source]
返回模型的总算力。
- 参数
as_string (bool, 可选) – 是否将算力作为字符串输出。默认为 False。
- 返回值
模型前向传播的乘加运算次数。
- get_total_macs(as_string=False)[source]
返回模型的总 MACs。
- 参数
as_string (bool, 可选) – 是否将算力作为字符串输出。默认为 False。
- 返回值
模型前向传播的乘加运算次数。
- get_total_duration(as_string=False)[source]
返回模型前向传播的总持续时间。
- 参数
as_string (bool, 可选) – 是否将持续时间作为字符串输出。默认为 False。
- 返回值
模型前向传播的延迟。
- get_total_params(as_string=False)[source]
返回每个排名存储的参数总数。
- 参数
as_string (bool, 可选) – 是否将参数作为字符串输出。默认为 False。
- 返回值
每个排名存储的参数总数。
- print_model_profile(profile_step=1, module_depth=-1, top_modules=1, detailed=True, output_file=None)[source]
打印模型图,其中每个模块都附加了测量的分析结果。
- 参数
profile_step (int, 可选) – 要分析的全局训练步骤。请注意,需要预热步骤才能获得准确的时间测量结果。
module_depth (int, 可选) – 要打印聚合模块信息的模型深度。设置为 -1 时,它会从顶部打印到最内部的模块(最大深度)。
top_modules (int, 可选) – 将聚合分析结果输出限制为指定的顶级模块数量。
detailed (bool, 可选) – 是否打印详细的模型分析结果。
output_file (str, 可选) – 输出文件的路径。如果为 None,则探查器会打印到标准输出。
- deepspeed.profiling.flops_profiler.profiler.get_model_profile(model, input_shape=None, args=[], kwargs={}, print_profile=True, detailed=True, module_depth=-1, top_modules=1, warm_up=1, as_string=True, output_file=None, ignore_modules=None, mode='forward')[source]
返回模型的总浮点运算次数、MACs 和参数。
示例
model = torchvision.models.alexnet() batch_size = 256 flops, macs, params = get_model_profile(model=model, input_shape=(batch_size, 3, 224, 224)))
- 参数
**model** ([torch.nn.Module]) – 要分析的 PyTorch 模型。
**input_shape** (tuple) – 模型的输入形状。如果指定,则模型将以该形状的张量作为唯一的 positional 参数。
**args** (list) – 模型的 positional 参数列表。
**kwargs** (dict) – 模型的关键字参数字典。
**print_profile** (bool, 可选) – 是否打印模型配置文件。默认为 True。
**detailed** (bool, 可选) – 是否打印详细的模型配置文件。默认为 True。
**module_depth** (int, 可选) – 嵌套模块的深度。默认为 -1(最内部的模块)。
**top_modules** (int, 可选) – 在聚合配置文件中打印的前几个模块的数量。默认为 3。
**warm_up** (int, 可选) – 在测量每个模块的延迟之前预热步骤的数量。默认为 1。
**as_string** (bool, 可选) – 是否将输出打印为字符串。默认为 True。
**output_file** (str, 可选) – 输出文件的路径。如果为 None,则分析器打印到标准输出。
**ignore_modules** ([type], 可选) – 在分析过程中要忽略的模块列表。默认为 None。
- 返回值
模型中的浮点运算次数、乘加运算次数 (MACs) 和参数数量。