How to get the output of the middle layer in the process of model training

[xiaohoua]

I know that I can use FX to get the output of the middle layer from this document，but it will ‘cut’ the model ,I don't think this is a good idea because the cut out part of the model also takes up memory.
And The pr2136 seems to solve a similar problem by forward_intermediates() API, but it only test model.eval().
Can you give me some advice about model.train()?（I want to build a loss function to fit the middle layer of the timm model and the middle layer of the open_clip model）
Looking forward to your reply！

[xiaohoua]

Now my approach is to directly modify the timm library source code：
in timm/models/vision_transformer.py line 418:

`

class VisionTransformer(nn.Module):
    
    dynamic_img_size: Final[bool]

    def __init__(
            self,
            **
    ) -> None:
    **self.atten_layer = {}**
        self.blocks = nn.Sequential(*[
            block_fn(
                dim=embed_dim,
                num_heads=num_heads,
                mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias,
                qk_norm=qk_norm,
                init_values=init_values,
                proj_drop=proj_drop_rate,
                attn_drop=attn_drop_rate,
                drop_path=dpr[i],
                norm_layer=norm_layer,
                act_layer=act_layer,
                mlp_layer=mlp_layer,
            )
            for i in range(depth)])
       def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.patch_embed(x)
        x = self._pos_embed(x)
        x = self.patch_drop(x)
        x = self.norm_pre(x)        
        if self.grad_checkpointing and not torch.jit.is_scripting():
            x = checkpoint_seq(self.blocks, x)
        else:
            # x = self.blocks(x)
            **for index,atten in enumerate(self.blocks):
                x = atten(x)
                if(index == 8 or index == 4):
                    self.atten_layer[index] = x**

        x = self.norm(x)
        return x     

`
** ** is the code I changed。
Then I can get the content I want through self.atten_layer, I don't know if this is correct, I think I should not modify the timm library source code, please give me some advice

[rwightman]

@xiaohoua it's fine to hack train scripts or models to do custom things, obviously it's not practical for the library itself to be flexible enough for all use cases, it would create too much overhead & mess for the common use cases.

That said, FX would probably work fine for this use case, it's probably possible to use forward hooks as well. forward_intermediates can work since you appear to want the intermediate block outputs and not something deeper within sub-blocks. In most cases you'd probably want to wrap the extraction bit w/ original model and whatever you need to do with the aux outputs in a parent module.

I have no idea what you're doing with it but random example below showing hypothetical forward of a wrapping module... forward_intermediates should work just fine in train model.

def forward(self, inpt):
  output, intermediates = self.model.forward_intermediates(inpt, indices=(4, 8), output_fmt='NLC')
  prediction = self.model.forward_head(output)
  aux_output = self.project(torch.cat(intermediates, -1))
  return prediction, aux_output 

[xiaohoua]

Thank you for your reply！
yes! I just want the intermediate block outputs. I am trying to use forward_intermediates but get an error:'CLIP' object has no attribute 'forward_intermediates'
My training code started out like this（No errors were reported during the training, but it took a long time）：

for epoch in range(start_epoch, args.num_epochs):
        epoch_loss = 0.

        for image, categories, embedding in tqdm.tqdm(iter(data_loader)):
            image = image.to(args.device, non_blocking=True)
            embedding = embedding.to(args.device, non_blocking=True)
            category = F.one_hot(categories, num_classes=397).float().to(args.device, non_blocking=True)
            ....................
            with autocast():  
                 ...............####this is my code#############
                output_embedding = student_model(image)
                student_atten_layer = student_model.atten_layer[8]
                student_atten_layer_adapted = adapter_student( student_atten_layer ）
                with torch.no_grad():
  
                    teacher_embedding = teacher_model.encode_image(image)
                    teacher_atten_layer = teacher_model.atten_layer[8]
                    teacher_atten_layer_adapted = adapter_teacher( teacher_atten_layer ）
                loss_attention = criterion_attention(teacher_atten_layer_adapted, student_atten_layer_adapted)
                loss_attention.backward()

I get an error when I use forward_intermediates

for epoch in range(start_epoch, args.num_epochs):
        epoch_loss = 0.

        for image, categories, embedding in tqdm.tqdm(iter(data_loader)):
            image = image.to(args.device, non_blocking=True)
            embedding = embedding.to(args.device, non_blocking=True)
            category = F.one_hot(categories, num_classes=397).float().to(args.device, non_blocking=True)
            ....................
            with autocast():  
                 ...............####use forward_intermediates###########
                output_embedding , student_atten_layer = teacher_model.forward_intermediates(image, indices=(4, 8), output_fmt='NLC')
                student_atten_layer = self.project(torch.cat(student_atten_layer , -1))
                student_atten_layer_adapted = adapter_teacher( student_atten_layer ）
                with torch.no_grad():
  
                    teacher_embedding, teacher_atten_layer  = teacher_model.forward_intermediates(image, indices=(4, 8), output_fmt='NLC')
                    teacher_atten_layer = self.project(torch.cat(teacher_atten_layer  , -1))
                    teacher_atten_layer_adapted = adapter_teacher( teacher_atten_layer ）
                loss_attention = criterion_attention(teacher_atten_layer_adapted, student_atten_layer_adapted)
                loss_attention.backward()

I don't know whether forward_intermediates is suitable for my needs.
Here, the teacher model only needs to obtain the middle layer output, not the final teacher_embedding, and the teacher model only needs to perform forward inference. However, the student model is required to complete the training process, the intermediate layer output and the final output_embedding are also required.
Teacher models are pretrained models created using open_clip.create_model_and_transforms, while student models are created using timm.create_model.

[xiaohoua]

Here's how to use hook to solve my problem for future reference，In this case I don't need to modify the open_clip library source code。

for epoch in range(start_epoch, args.num_epochs):
        epoch_loss = 0.
        hook_handle, activations = get_activation('atten_layer_8')
        hook = teacher_model.visual.transformer.resblocks[8].register_forward_hook(hook_handle)
        for image, categories, embedding in tqdm.tqdm(iter(data_loader)):
            image = image.to(args.device, non_blocking=True)
            embedding = embedding.to(args.device, non_blocking=True)
            category = F.one_hot(categories, num_classes=397).float().to(args.device, non_blocking=True)
            ....................
            with autocast():  
                 ...............
                output_embedding , student_atten_layer = teacher_model.forward_intermediates(image, indices=(4, 8), output_fmt='NLC')
                student_atten_layer = self.project(torch.cat(student_atten_layer , -1))
                student_atten_layer_adapted = adapter_teacher( student_atten_layer ）
                with torch.no_grad():
                   ####use hook###########
                    _ = teacher_model.encode_image(image_teacher)
                    teacher_atten_layer = activations['atten_layer_8']
                    teacher_atten_layer_adapted = adapter_teacher( teacher_atten_layer ）
                loss_attention = criterion_attention(teacher_atten_layer_adapted, student_atten_layer_adapted)
                loss_attention.backward()