Adding print statements to VisionTransformer

[NielsRogge]

Hi,

I'd like to understand the ViT model better by adding some print statements to the forward pass of ViT, to check the intermediate hidden states, verify shapes and so on.

I've forked pytorch-image-models and added some print statements as follows:

def forward_features(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        x = self.pos_drop(x)

        print("Shape of hidden states before encoder layers:")
        print(x.shape)

        for idx, blk in enumerate(self.blocks):
            print(f"Hidden states before layer {idx}")
            print(x[0,:3,:3])

            x = blk(x)

            print(f"Hidden states after layer {idx}")
            print(x[0,:3,:3])

        x = self.norm(x)[:, 0]
        x = self.pre_logits(x)
        return x

Are the hidden states of shape (batch_size, seq_len, hidden_size) as in HuggingFace Transformers?
However, when using this to perform a forward pass on an image, nothing gets printed. Any idea why this is the case?
Notebook:
https://colab.research.google.com/drive/1RQBYatwWApnEAPyFG3lNvci0hj3vZxIv?usp=sharing

[rwightman]

I litter my models with prints all the time when debugging and getting them working. Maybe you clobbered over your modified copy of pytorch-image-models with the pip timm?

[NielsRogge]

Thanks for your reply!

Indeed, fixed it by instead using pip install git+https://github.com/NielsRogge/pytorch-image-models.git@understanding_vit.