Regression on the vision transformers

1

I am trying to do regression on the vit_b_16 models to estimate values from the image, but I get the issue when I try to replace the last classification layer with regression layer,

class RegressionViT(nn.Module):
def init(self, in_features=224 * 224 * 3, num_classes=1, pretrained=True):
super(RegressionViT, self).init()
self.vit_b_16 = vit_b_16(pretrained=pretrained) # Load pre-trained weights

    # Replace the final classification layer with a regression head
    self.regressor = nn.Linear(self.vit_b_16.heads.in_features, num_classes)

def forward(self, x):
    x = self.vit_b_16(x)
    x = self.regressor(x)
    return x

I get the error for the following code

AttributeError: ‘VisionTransformer’ object has no attribute ‘head’

What will be the correct code to do the regression?

2

Your code fails with a different error message for me:

# AttributeError: 'Sequential' object has no attribute 'in_features'

After fixing this by indexing the self.vit_b_16.heads[0] container, the code works.