如何打印预训练TimeSformer模型各层的输入输出尺寸？

问题背景 #

在自定义nn.Module子类中，将预训练TimeSformer模型作为self.feature属性，尝试打印所有层的输入输出尺寸时，遇到TypeError: 'TimeSformer' object is not iterable错误，后续使用children()的尝试也未达到预期效果。

模型结构概述 #

TimeSformer(
  (model): VisionTransformer(
    (dropout): Dropout(p=0.0, inplace=False)
    (patch_embed): PatchEmbed(
      (proj): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16))
    )
    (pos_drop): Dropout(p=0.0, inplace=False)
    (time_drop): Dropout(p=0.0, inplace=False)
    (blocks): ModuleList(  
      (0): Block(...)  # 共12个Block，每个包含空间注意力、时间注意力、MLP等子模块
      ...
      (11): Block(...)
    )
    (norm): LayerNorm((768,), eps=1e-06, elementwise_affine=True) 
    (head): Linear(in_features=768, out_features=400, bias=True) 
  )
)

错误代码示例 #

最初尝试直接遍历self.feature导致报错：

class Class(nn.Module):
    def __init__(self, pretrained=False):
        super(Class, self).__init__()
        
        self.feature = TimeSformer(img_size=224, num_classes=400, num_frames=8, attention_type='divided_space_time',  
                                   pretrained_model='path/to/the/weight.pyth')

    def forward(self, x):
        for layer in self.feature:
            x = layer(x)
            print(x.size())
        return x

报错信息：

TypeError: 'TimeSformer' object is not iterable

可行解决方案 #

方案1：利用前向钩子自动记录所有层IO #

通过注册前向钩子，递归捕获所有子模块的输入输出尺寸，无需手动拆解模型：

import torch
import torch.nn as nn
from timesformer.models.vit import TimeSformer

class Class(nn.Module):
    def __init__(self, pretrained=False):
        super(Class, self).__init__()
        self.feature = TimeSformer(img_size=224, num_classes=400, num_frames=8, attention_type='divided_space_time',  
                                   pretrained_model='path/to/the/weight.pyth')

    def _log_io(self, module, input, output):
        # 打印模块类名、输入尺寸、输出尺寸
        input_shape = tuple(input[0].shape) if input else "无输入"
        output_shape = tuple(output.shape) if output is not None else "无输出"
        print(f"[{module.__class__.__name__}] 输入尺寸: {input_shape} | 输出尺寸: {output_shape}")

    def forward(self, x):
        # 给所有子模块注册钩子
        hooks = []
        # 注意要访问TimeSformer内部的model属性（VisionTransformer核心）
        for sub_module in self.feature.model.modules():
            hook = sub_module.register_forward_hook(self._log_io)
            hooks.append(hook)
        
        # 执行前向传播
        output = self.feature(x)
        
        # 移除钩子，避免重复注册
        for hook in hooks:
            hook.remove()
        
        return output

方案2：手动逐层遍历（针对TimeSformer结构定制） #

如果需要更清晰的层级展示，可根据模型结构手动拆解前向流程，逐个打印尺寸：

import torch
import torch.nn as nn
from timesformer.models.vit import TimeSformer

class Class(nn.Module):
    def __init__(self, pretrained=False):
        super(Class, self).__init__()
        self.feature = TimeSformer(img_size=224, num_classes=400, num_frames=8, attention_type='divided_space_time',  
                                   pretrained_model='path/to/the/weight.pyth')
        # 提取核心VisionTransformer模块
        self.vit = self.feature.model

    def forward(self, x):
        print(f"初始输入尺寸: {x.shape}")
        
        # Patch嵌入层
        x = self.vit.patch_embed(x)
        print(f"PatchEmbed输出尺寸: {x.shape}")
        
        # 位置Dropout
        x = self.vit.pos_drop(x)
        print(f"pos_drop输出尺寸: {x.shape}")
        
        # 时间Dropout
        x = self.vit.time_drop(x)
        print(f"time_drop输出尺寸: {x.shape}")
        
        # 遍历所有Block
        for idx, block in enumerate(self.vit.blocks):
            print(f"\n=== Block {idx} ===")
            # 空间注意力分支
            norm1_out = block.norm1(x)
            print(f"norm1输出尺寸: {norm1_out.shape}")
            attn_out = block.attn(norm1_out)
            print(f"空间attn输出尺寸: {attn_out.shape}")
            x = x + block.drop_path(attn_out)
            
            # 时间注意力分支
            temp_norm_out = block.temporal_norm1(x)
            print(f"temporal_norm1输出尺寸: {temp_norm_out.shape}")
            temp_attn_out = block.temporal_attn(temp_norm_out)
            print(f"时间attn输出尺寸: {temp_attn_out.shape}")
            temp_fc_out = block.temporal_fc(temp_attn_out)
            print(f"temporal_fc输出尺寸: {temp_fc_out.shape}")
            x = x + block.drop_path(temp_fc_out)
            
            # MLP分支
            norm2_out = block.norm2(x)
            print(f"norm2输出尺寸: {norm2_out.shape}")
            mlp_out = block.mlp(norm2_out)
            print(f"MLP输出尺寸: {mlp_out.shape}")
            x = x + block.drop_path(mlp_out)
            print(f"Block {idx} 最终输出尺寸: {x.shape}")
        
        # 最终LayerNorm
        x = self.vit.norm(x)
        print(f"\n最终Norm输出尺寸: {x.shape}")
        
        # 分类头
        x = self.vit.head(x)
        print(f"分类头输出尺寸: {x.shape}")
        
        return x

关键注意点 #

• TimeSformer模块的核心逻辑在内部的model属性（即VisionTransformer）中，不能直接遍历self.feature，需访问self.feature.model。
• modules()方法会递归获取所有嵌套子模块，适合用钩子自动记录；children()仅获取直接子模块，无法覆盖深层嵌套的子层。

内容的提问来源于stack exchange，提问作者dtr43