why model.fc.out_features not matching the logits dimensions? (pytorch.nn)

Viewed 139

just implemented resnet50, fined tuned it with output features=200, then loaded the best epoch of the trained models. Then when I predicted the logits, I'm getting the wrong dimension of 1000. How I can change it to 200?

The data is tiny-imageNet:

data_dir = 'tiny-224/'
num_workers = {'train' : 6,'val' : 0,'test'  : 0}
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomRotation(20),
        transforms.RandomHorizontalFlip(0.5),
        transforms.ToTensor(),
        transforms.Normalize([0.4802, 0.4481, 0.3975], [0.2302, 0.2265, 0.2262]),
    ]),
    'val': transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.4802, 0.4481, 0.3975], [0.2302, 0.2265, 0.2262]),
    ]),
    'test': transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.4802, 0.4481, 0.3975], [0.2302, 0.2265, 0.2262]),
    ])
}
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
                  for x in ['train', 'val','test']}
dataloaders = {x: data.DataLoader(image_datasets[x], batch_size=100, shuffle=True, num_workers=num_workers[x])
                  for x in ['train', 'val', 'test']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val', 'test']}

Loading model:

model_ft = models.resnet50(True)
#Finetune Final few layers to adjust for tiny imagenet input
model_ft.avgpool = nn.AdaptiveAvgPool2d(1)
model_ft.fc.out_features = 200
PATH = r'./models/224/model_12_epoch.pt'
model_ft.load_state_dict(torch.load(PATH))

now when I run the prediction it returns the wrong dimensions :

images, labels = next(iter(dataloaders['val']))
logits = model_ft(images)

logits dimension : (100,1000) when expected (100,200)

1 Answers

You can try this:

model_ft = models.resnet50(True)
model_ft.avgpool = nn.AdaptiveAvgPool2d(1)
num_ftrs = model_ft.fc.in_features        # read the input of the last layer
model_ft.fc = nn.Linear(num_ftrs, 200)    # replace the last layer with linear layer where
                                          # input = num_ftrs and output = 200
Related