Extracting certain blocks from a pretrained model and removing some layers

Tm4: <p>I want to extract the features from certain blocks of the <a href="https://github.com/facebookresearch/TimeSformer" rel="noopener nofollow ugc">TimeSformer</a> model and also want to remove the last two layers.</p>
<pre><code class="lang-auto">
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')
</code></pre>
<p><strong>The print of the model is as follows:</strong></p>
<p>TimeSformer(<br>
(model): VisionTransformer(<br>
(dropout): Dropout(p=0.0, inplace=False)<br>
(patch_embed): PatchEmbed(<br>
(proj): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16))<br>
)<br>
(pos_drop): Dropout(p=0.0, inplace=False)<br>
(time_drop): Dropout(p=0.0, inplace=False)<br>
(blocks): ModuleList(<br>
(0): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(temporal_norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(temporal_attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(temporal_fc): Linear(in_features=768, out_features=768, bias=True)<br>
(drop_path): Identity()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(1): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(temporal_norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(temporal_attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(temporal_fc): Linear(in_features=768, out_features=768, bias=True)<br>
(drop_path): DropPath()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
(2): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(temporal_norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(temporal_attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(temporal_fc): Linear(in_features=768, out_features=768, bias=True)<br>
(drop_path): DropPath()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
.<br>
.<br>
.<br>
.<br>
.<br>
(11): Block(<br>
(norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(temporal_norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(temporal_attn): Attention(<br>
(qkv): Linear(in_features=768, out_features=2304, bias=True)<br>
(proj): Linear(in_features=768, out_features=768, bias=True)<br>
(proj_drop): Dropout(p=0.0, inplace=False)<br>
(attn_drop): Dropout(p=0.0, inplace=False)<br>
)<br>
(temporal_fc): Linear(in_features=768, out_features=768, bias=True)<br>
(drop_path): DropPath()<br>
(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)<br>
(mlp): Mlp(<br>
(fc1): Linear(in_features=768, out_features=3072, bias=True)<br>
(act): GELU()<br>
(fc2): Linear(in_features=3072, out_features=768, bias=True)<br>
(drop): Dropout(p=0.0, inplace=False)<br>
)<br>
)<br>
)<br>
(norm): LayerNorm((768,), eps=1e-06, elementwise_affine=True) **** I want to remove this layer*****<br>
(head): Linear(in_features=768, out_features=400, bias=True)   **** I want to remove this layer*****<br>
)<br>
)</p>
<p>Specifically, I want to extract the outputs of the 4th, 8th and 11th layers of the model and removing the lats two layers. How can I do this. I tried using  <code>TimeSformer.blocks[0]</code> but that was not working.</p>