浮点运算量分析器

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, optional) – 分析时要忽略的模块列表。默认为 None。

stop_profile()[source]

停止分析。

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

reset_profile()[source]

重置分析。

添加或重置额外属性。

end_profile()[source]

结束分析。

添加的属性和句柄从所有模块上递归地移除。

get_total_flops(as_string=False)[source]

返回模型的总浮点运算量。

参数: as_string (bool, optional) – 是否以字符串形式输出浮点运算量。默认为 False。
返回值: 模型正向传播的乘加运算次数。

get_total_macs(as_string=False)[source]

返回模型的总 MAC 次数。

参数: as_string (bool, optional) – 是否以字符串形式输出浮点运算量。默认为 False。
返回值: 模型正向传播的乘加运算次数。

get_total_duration(as_string=False)[source]

返回模型正向传播的总持续时间。

参数: as_string (bool, optional) – 是否以字符串形式输出持续时间。默认为 False。
返回值: 模型正向传播的延迟。

get_total_params(as_string=False)[source]

返回每个排名存储的参数总数。

参数: as_string (bool, optional) – 是否以字符串形式输出参数。默认为 False。
返回值: 每个排名存储的参数总数。

print_model_profile(profile_step=1, module_depth=-1, top_modules=1, detailed=True, output_file=None)[source]

打印带有测量分析的模型图，该分析附加到每个模块。

参数

profile_step (int, optional) – 分析的全局训练步骤。请注意，需要预热步骤才能获得准确的时间测量结果。
module_depth (int, optional) – 打印汇总模块信息的模型深度。当设置为 -1 时，它会从顶部到最内层的模块打印信息（最大深度）。
top_modules (int, optional) – 将汇总分析输出限制为指定的 top 模块数量。
detailed (bool, optional) – 是否打印详细的模型分析。
output_file (str, optional) – 输出文件的路径。如果为 None，则分析器会打印到标准输出。

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

打印深度为 module_depth 处按汇总时间、浮点运算和参数排序的 top_modules 个 top 模块的名称。

参数

module_depth (int, optional) – 要显示的模块深度。默认为 -1（最内层的模块）。
top_modules (int, optional) – 要显示的 top 模块数量。默认为 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]

返回模型的总浮点运算次数、MAC 和参数。

示例

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) – 模型的输入形状。如果指定，则模型将以此形状的张量作为唯一的定位参数。
args (list) – 模型的定位参数列表。
kwargs (dict) – 模型的关键字参数字典。
print_profile (bool, optional) – 是否打印模型分析结果。默认值为 True。
detailed (bool, optional) – 是否打印详细的模型分析结果。默认值为 True。
module_depth (int, optional) – 嵌套模块的深度。默认值为 -1（最内层的模块）。
top_modules (int, optional) – 在聚合分析结果中打印的顶级模块数量。默认值为 3。
warm_up (int, optional) – 在测量每个模块的延迟之前预热的步数。默认值为 1。
as_string (bool, optional) – 是否以字符串形式打印输出。默认值为 True。
output_file (str, optional) – 输出文件的路径。如果为 None，则分析器将打印到标准输出。
ignore_modules ([type], optional) – 分析过程中要忽略的模块列表。默认值为 None。

返回值

模型中的浮点运算次数、乘加运算次数（MAC）和参数的数量。