算力探查器


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。

stop_profile()[source]

停止分析。

所有 torch.nn.functionals 都会恢复到其原始状态。

reset_profile()[source]

重置分析。

添加或重置额外属性。

end_profile()[source]

结束分析。

会递归地从所有模块中删除添加的属性和句柄。

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,则探查器会打印到标准输出。

print_model_aggregated_profile(module_depth=-1, top_modules=1)[source]

打印根据聚合时间、算力和参数在深度 module_depth 处的前 top_modules 个模块的名称。

参数
  • module_depth (int, 可选) – 要显示的模块深度。默认为 -1(最内部的模块)。

  • top_modules (int, 可选) – 要显示的顶级模块数量。默认为 1。

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) 和参数数量。