在 PyTorch 中访问预训练模型中的特定层答案

【问题标题】：Accessing a specific layer in a pretrained model in PyTorch在 PyTorch 中访问预训练模型中的特定层
【发布时间】：2022-08-10 17:17:34
【问题描述】：

我想从TimeSformer 模型的某些块中提取特征，并且还想删除最后两层。

import torch
from timesformer.models.vit import TimeSformer

model = TimeSformer(img_size=224, num_classes=400, num_frames=8, attention_type=\'divided_space_time\',  pretrained_model=\'/path/to/pretrained/model.pyth\')

模型打印如下：

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(
    (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
    (attn): Attention(
      (qkv): Linear(in_features=768, out_features=2304, bias=True)
      (proj): Linear(in_features=768, out_features=768, bias=True)
      (proj_drop): Dropout(p=0.0, inplace=False)
      (attn_drop): Dropout(p=0.0, inplace=False)
    )
    (temporal_norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
    (temporal_attn): Attention(
      (qkv): Linear(in_features=768, out_features=2304, bias=True)
      (proj): Linear(in_features=768, out_features=768, bias=True)
      (proj_drop): Dropout(p=0.0, inplace=False)
      (attn_drop): Dropout(p=0.0, inplace=False)
    )
    (temporal_fc): Linear(in_features=768, out_features=768, bias=True)
    (drop_path): Identity()
    (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
    (mlp): Mlp(
      (fc1): Linear(in_features=768, out_features=3072, bias=True)
      (act): GELU()
      (fc2): Linear(in_features=3072, out_features=768, bias=True)
      (drop): Dropout(p=0.0, inplace=False)
    )
  )
  (1): Block(
    (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
    (attn): Attention(
      (qkv): Linear(in_features=768, out_features=2304, bias=True)
      (proj): Linear(in_features=768, out_features=768, bias=True)
      (proj_drop): Dropout(p=0.0, inplace=False)
      (attn_drop): Dropout(p=0.0, inplace=False)
    )
    (temporal_norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
    (temporal_attn): Attention(
      (qkv): Linear(in_features=768, out_features=2304, bias=True)
      (proj): Linear(in_features=768, out_features=768, bias=True)
      (proj_drop): Dropout(p=0.0, inplace=False)
      (attn_drop): Dropout(p=0.0, inplace=False)
    )
    (temporal_fc): Linear(in_features=768, out_features=768, bias=True)
    (drop_path): DropPath()
    (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
    (mlp): Mlp(
      (fc1): Linear(in_features=768, out_features=3072, bias=True)
      (act): GELU()
      (fc2): Linear(in_features=3072, out_features=768, bias=True)
      (drop): Dropout(p=0.0, inplace=False)
    )
  )
.
.
.
.
.
.
  (11): Block(
    (norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
    (attn): Attention(
      (qkv): Linear(in_features=768, out_features=2304, bias=True)
      (proj): Linear(in_features=768, out_features=768, bias=True)
      (proj_drop): Dropout(p=0.0, inplace=False)
      (attn_drop): Dropout(p=0.0, inplace=False)
    )
    (temporal_norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
    (temporal_attn): Attention(
      (qkv): Linear(in_features=768, out_features=2304, bias=True)
      (proj): Linear(in_features=768, out_features=768, bias=True)
      (proj_drop): Dropout(p=0.0, inplace=False)
      (attn_drop): Dropout(p=0.0, inplace=False)
    )
    (temporal_fc): Linear(in_features=768, out_features=768, bias=True)
    (drop_path): DropPath()
    (norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)
    (mlp): Mlp(
      (fc1): Linear(in_features=768, out_features=3072, bias=True)
      (act): GELU()
      (fc2): Linear(in_features=3072, out_features=768, bias=True)
      (drop): Dropout(p=0.0, inplace=False)
    )
  )
)
(norm): LayerNorm((768,), eps=1e-06, elementwise_affine=True) **** I want to remove this layer*****
(head): Linear(in_features=768, out_features=400, bias=True) **** I want to remove this layer*****

)

具体来说，我想提取模型的第 4、8 和 11 个块的输出，并删除 lats 两层。我怎样才能做到这一点。我尝试使用 TimeSformer.blocks[0] 但这不起作用。

更新：

我有一个类，我需要访问上述 TimeSformer 块作为该类的输出。这个类的输入是一个 5D 张量。这是我用于提取上述块的输出的未修改代码：

class Model(nn.Module):
def __init__(self, pretrained=False):
    super(Model, self).__init__()
    
    
    self.model =TimeSformer(img_size=224, num_classes=400, num_frames=8, attention_type=\'divided_space_time\',  
                                       pretrained_model=\'/home/user/models/TimeSformer_divST_16x16_448_K400.pyth\')
    
   
    self.activation = {}
    def get_activation(name):
        def hook(model, input, output):
            self.activation[name] = output.detach()
            return hook

    self.model.model.blocks[4].register_forward_hook(get_activation(\'block4\'))
    self.model.model.blocks[8].register_forward_hook(get_activation(\'block8\'))
    self.model.model.blocks[11].register_forward_hook(get_activation(\'block11\'))


    block4_output = self.activation[\'block4\']
    block8_output = self.activation[\'block8\']
    block11_output = self.activation[\'block11\']
    
    
    
def forward(self, x, out_consp = False):
    
    features2, features3, features4 = self.model(x)

标签： python deep-learning pytorch pre-trained-model

【解决方案1】：

要从特定层提取中间输出，可以将其注册为钩子，示例如下所示：

import torch
from timesformer.models.vit import TimeSformer

model = TimeSformer(img_size=224, num_classes=400, num_frames=8, attention_type='divided_space_time',  pretrained_model='/path/to/pretrained/model.pyth')

activation = {}
def get_activation(name):
    def hook(model, input, output):
        activation[name] = output.detach()
    return hook

model.model.blocks[4].register_forward_hook(get_activation('block4'))
model.model.blocks[8].register_forward_hook(get_activation('block8'))
model.model.blocks[11].register_forward_hook(get_activation('block11'))

x = torch.randn(3,3,224,224)
output = model(x)

block4_output = activation['block4']
block8_output = activation['block8']
block11_output = activation['block11']

要删除最后两层，您可以将它们替换为 Identity：

model.norm = torch.nn.Identity()
model.head= torch.nn.Identity()

【讨论】：

感谢您的回答，但我面临此错误：未定义名称“激活”。此外，我需要提取它们的输出的块位于 ModuleList 中，我用 ****** 突出显示了 ModuleList。我认为由于模型的结构，代码无法访问这些块。与层是顺序的 resnet 模型相比，TimeSformer 模型块位于 ModuleList 内。
如果您尝试在课堂上初始化它，请将activation 设置为self.activation 否则我无法理解为什么会出现该错误。 ModuleList 也是可以交互的，只要调用块的索引就可以了，你可以编辑问题并上传你的完整代码
我添加了我想在其中使用您的答案的代码。抱歉，代码的总量如此之大。作为一个问题，在需要的情况下，如何访问例如 1st Block 的 (proj) 层。