模块工具

class torcheval.tools.ModuleSummary

模块及其子模块的摘要。它收集以下信息

• 名称

• 类型

• 参数数量

• 可训练参数的数量

• 估计大小（字节）

• 该模块是否包含未初始化的参数

• 正向计算的 FLOPs（“？”表示未计算）

• 反向计算的 FLOPs（“？”表示未计算）

• 输入形状（“？”表示未计算）

• 输出形状（“？”表示未计算）

• 正向计算的经过时间（毫秒）（“？”表示未计算）

property flops_backward: int | Literal['?']

返回使用该模块进行反向计算的总 FLOPs。

property flops_forward: int | Literal['?']

返回使用该模块进行正向计算的总 FLOPs。

property forward_elapsed_time_ms: Literal['?'] | float

返回模块的正向计算时间（毫秒）。

property has_uninitialized_param: bool

返回该模块中是否存在未初始化的参数

property in_size: Literal['?'] | List[int]

返回模块的输入大小

property module_name: str

返回该模块的名称

property module_type: str

返回该模块的类型。

property num_parameters: int

返回该模块中参数的总数。

property num_trainable_parameters: int

返回该模块中可训练参数（requires_grad=True）的总数。

property out_size: Literal['?'] | List[int]

返回模块的输出大小

property size_bytes: int

返回模块的总估计大小（字节）。

property submodule_summaries: Dict[str, ModuleSummary]

一个字典，以子模块的名称作为键，对应 ModuleSummary 对象作为值。这些可以用于遍历和可视化。

torcheval.tools.get_module_summary(module: Module, module_args: Tuple[Any, ...] | None = None, module_kwargs: MutableMapping[str, Any] | None = None) → ModuleSummary

生成一个 ModuleSummary 对象，然后赋值并生成子模块树。

参数：

• module – 要总结的模块。

• module_args – 模块运行并计算 FLOPs 和激活大小的实参元组。

• module_kwargs –

  要传递给模块正向函数的任何关键字实参。

  注意

  要计算 FLOPs，您必须使用 PyTorch 1.13 或更高版本。

  注意

  如果模块包含任何延迟子模块，我们将不会计算 FLOPs。

  注意

  目前仅支持输出单个张量的模块。TODO：支持更灵活的模块输出。

torcheval.tools.get_summary_table(module_summary: ModuleSummary, human_readable_nums: bool = True) → str

生成一个字符串 summary_table，将 module_summary 中的信息制成表格。

参数：

• module_summary – 要打印/制成表格的 module_summary

• human_readable_nums – 如果设置为 False，则使用精确值（例如 1234 vs 1.2 K）

torcheval.tools.prune_module_summary(module_summary: ModuleSummary, *, max_depth: int) → None

修剪模块摘要树中深度超过 max_depth 的模块摘要。ModuleSummary 对象将在原位进行修剪。

参数：

• module_summary – 要修剪的根模块摘要。

• max_depth – 要保留的模块摘要的最大深度。

Raises:

ValueError – 如果 max_depth 是小于 1 的整数