模块工具

class torcheval.tools.ModuleSummary

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

• 名称

• 类型

• 参数数量

• 可训练参数的数量

• 估计的字节大小

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

• 前向传播的FLOPs（“？”表示未计算）

• 反向传播的FLOPs（“?”表示未计算）

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

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

• 以毫秒为单位的前进耗时（“?”表示未计算）

property flops_backward: int | Literal['?']

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

property flops_forward: int | Literal['?']

返回使用此模块进行前向计算的总浮点运算次数。

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对象作为值。这些可以遍历以进行可视化。

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

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

Parameters:

• module – 要总结的模块。

• module_args – 一个包含模块运行和计算FLOPs和激活大小的参数的元组。

• module_kwargs –

  任何要传递给模块的前向函数的kwargs参数。

  注意

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

  注意

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

  注意

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

torcheval.tools.get_summary_table(module_summary: 模块摘要, human_readable_nums: bool = True) → str

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

Parameters:

• module_summary – 要打印/表格化的module_summary

• human_readable_nums – 设置为False以获得精确值（例如1234 vs 1.2 K）

torcheval.tools.prune_module_summary(module_summary: 模块摘要, *, max_depth: int) → None

修剪模块摘要树中深度超过max_depth的模块摘要。ModuleSummary对象会被就地修剪。

Parameters:

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

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

Raises:

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