PyTorch获取模型输入输出Shape详解

"在PyTorch中，与TensorFlow或Caffe不同，官方并没有提供直接获取模型input/output shape的功能。然而，可以通过编写自定义代码来实现这一目的。本实例提供了一种解决办法，但需要注意，由于不同的模块如CNN、RNN实现方式不同，可能需要针对特定模块进行代码调整。这个方法依赖于对模型进行一次前向传播（model(x)），以计算出形状信息。" 在PyTorch中，获取模型输入和输出的形状通常是通过动态图的特性实现的，而不是像TensorFlow或Caffe那样有内置的模型摘要工具。在给定的代码片段中，我们看到一个工作示例，它使用了`OrderedDict`来存储每个层的输入和输出形状。首先，导入所需的库，包括`collections.OrderedDict`、`torch`、`torch.autograd.Variable`、`torch.nn`以及一个假设存在的`models.crnn`模块。 `get_output_size`函数是一个递归函数，用于处理模型输出可能是单个张量或元组的情况。如果输出是元组（多输出模型），它会遍历每个元素并递归地获取它们的形状。对于单个张量，它简单地将形状转换为列表并存储到`summary_dict`中。 `summary`函数是核心部分，它接收输入尺寸`input_size`和模型对象`model`作为参数。通过注册钩子函数`hook`到模型的每个模块，我们可以在前向传播过程中捕获输入和输出。`hook`函数记录了模块的类名（class_name）、模块的索引（module_idx）以及输入和输出的形状。 `register_hook`函数遍历模型的子模块，并对每个子模块调用`hook`函数进行注册。这样，当模型被调用时，每个模块的输入和输出都会被记录下来。 请注意，这种方法依赖于给定一个输入尺寸`input_size`来执行前向传播，因此获取的形状信息是基于这个特定输入的。在实际应用中，这可能需要根据训练数据的形状进行调整。 这段代码提供了一个实用的解决方案，尽管它可能不适用于所有类型的PyTorch模块，特别是那些具有特殊特性的模块，如RNN，可能需要对代码进行调整以正确处理它们的权重和偏置。然而，对于大多数标准的卷积神经网络和全连接网络，这个方法应该是足够的。