Search code examples
python-3.xtorchfeature-extractionvision-transformer

Vision transformer: Visualize feature maps


I am working on visualizing feature maps of my vision transformer but i am unable to visualize feature maps. When i print model.children() it shows convolution layers but still i cannot verify the if statement.

list(model.children())

Output

[OverlapPatchEmbed(
   (proj): Conv2d(3, 64, kernel_size=(7, 7), stride=(4, 4), padding=(3, 3))
   (norm): LayerNorm((64,), eps=1e-05, elementwise_affine=True)
 ),
 OverlapPatchEmbed(
   (proj): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
   (norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
 ),
 OverlapPatchEmbed(
   (proj): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
   (norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
 ),
 OverlapPatchEmbed(
   (proj): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
   (norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)
 ),
 ModuleList(
   (0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=64)
   (1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=64)
   (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=64)
   (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), paddin...

I want to access Conv2d and visulize the feature map but i am unable to do so type(model_children[i]) == Conv2d is not True and i have no idea why?

model_children = list(model.children())
# counter to keep count of the conv layers
counter = 0
# append all the conv layers and their respective wights to the list
for i in range(len(model_children)):
    if type(model_children[i]) == Conv2d:
        counter += 1
        model_weights.append(model_children[i].weight)
        conv_layers.append(model_children[i])
    elif type(model_children[i]) == nn.Sequential:
        for j in range(len(model_children[i])):
            for child in model_children[i][j].children():
                if type(child) == nn.Conv2d:
                    counter += 1
                    model_weights.append(child.weight)
                    conv_layers.append(child)
print(f"Total convolution layers: {counter}")
print("conv_layers")

Solution

  • Actually model_children[i].weight does not contain weight attribute. Inside OverlapPatchEmbed, the proj layer contains Conv2d and Conv2d consists of weight attribute. You can correct it below.

    if model_children[i] == model.patch_embed1:
        counter += 1
        weigh = model_children[i].proj
        model_weights.append(weigh.weight)
        conv_layers.append(model_children[i].proj)
    elif model_children[i] == model.patch_embed2:
        counter += 1
        weigh = model_children[i].proj
        model_weights.append(weigh.weight)
        conv_layers.append(model_children[i].proj)