pytorch中获取模型input/output shape实例

在PyTorch中，获取模型的输入(input)和输出(output)形状(shape)并不像在TensorFlow或Caffe那样直接，因为PyTorch的设计更注重灵活性。然而，可以通过编写自定义代码来实现这一功能。以下是一个实例，展示了如何通过遍历模型的层并计算其前向传播输出的形状来获取输入和输出形状。 我们需要导入必要的库： ```python from collections import OrderedDict import torch from torch.autograd import Variable import torch.nn as nn ``` 接下来，定义一个函数`get_output_size`，这个函数递归地处理输出，如果输出是元组，它会继续处理每个元素： ```python def get_output_size(summary_dict, output): if isinstance(output, tuple): for i in range(len(output)): summary_dict[i] = OrderedDict() summary_dict[i] = get_output_size(summary_dict[i], output[i]) else: summary_dict['output_shape'] = list(output.size()) return summary_dict ``` 然后，定义主函数`summary`，它接受输入尺寸(input_size)和模型(model)作为参数。这个函数会遍历模型的所有层，并记录输入和输出形状： ```python def summary(input_size, model): def register_hook(module): def hook(module, input, output): class_name = str(module.__class__).split('.')[-1].split("'")[0] module_idx = len(summary) m_key = '%s-%i' % (class_name, module_idx + 1) summary[m_key] = OrderedDict() summary[m_key]['input_shape'] = list(input[0].size()) summary[m_key] = get_output_size(summary[m_key], output) params = 0 if hasattr(module, 'weight'): params += torch.prod(torch.LongTensor(list(module.weight.size()))) if module.weight.requires_grad: summary[m_key]['trainable'] = True else: summary[m_key]['trainable'] = False # 如果有偏置项，可以添加类似处理 # if hasattr(module, 'bias'): # params += torch.prod(torch.LongTensor(list(module.bias.size()))) summary[m_key]['nb_params'] = params if not isinstance(module, nn.Sequential) and \ not isinstance(module, nn.ModuleList) and \ not (module == model): hooks.append(module.register_forward_hook(hook)) # 检查是否有多个输入到网络 if isinstance(input_size[0], (list, tuple)): x = [Variable(torch.rand(1, *in_size)) for in_size in input_size] else: x = Variable(torch.rand(1, *input_size)) # 创建属性 summary = OrderedDict() hooks = [] # 注册hook model.apply(register_hook) # 运行前向传播 with torch.no_grad(): model(x) # 移除所有hook for h in hooks: h.remove() return summary ``` 在上述代码中，`register_hook`是一个内部辅助函数，用于注册一个hook到每个模块上，当前向传播执行时，`hook`函数会被调用，从而记录每个模块的输入和输出形状。`summary`函数首先创建一个空的OrderedDict，然后遍历模型的所有子模块，注册hook。在运行前向传播后，所有的形状信息都会被记录下来。 注意，这个方法需要构造一个实际的输入数据调用`model(x)`，以便让模型的前向传播执行。此外，此方法仅适用于那些权重存储在`weight`属性中的模块，对于像RNN这样的模块，可能需要额外的处理，因为它们的权重和偏置存储方式可能不同。 你可以通过调用`summary(input_size, model)`并传入你的模型和输入尺寸来获取模型的输入输出形状信息。例如，如果你有一个卷积神经网络(CNN)，并且知道输入图像的尺寸是224x224x3，你可以这样使用： ```python input_size = (3, 224, 224) model = your_cnn_model print(summary(input_size, model)) ``` 这将输出一个详细的字典，包含每个层的输入形状、输出形状、是否训练以及参数数量。这个信息对于理解和调试模型非常有用。