I’m a rookie in PyTorch.
I have modified the code of tutorial. (Transfer Learning tutorial 115 )
Finetuning the convnet
dset_classes_number = len(dset_classes)
model_ft = models.alexnet(pretrained=True) 
model_ft.classifier  = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, dset_classes_number),
)
if use_gpu:
    model_ft = model_ft.cuda()
criterion = nn.CrossEntropyLoss()
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
# Train and evaluate
model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                       num_epochs=25)

ConvNet as fixed feature extractor
dset_classes_number = len(dset_classes)
model_conv = torchvision.models.alexnet(pretrained=True)
for param in model_conv.parameters():
    param.requires_grad = False

# num_ftrs = model_conv.fc.in_features
# model_conv.fc = nn.Linear(num_ftrs, dset_classes_number)

model_conv.classifier  = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, dset_classes_number),
)
if use_gpu:
    model_conv = model_conv.cuda()

criterion = nn.CrossEntropyLoss()
optimizer_conv = optim.SGD(model_conv.classifier.parameters(), lr=0.001, momentum=0.9)
# Train and evaluate
model_conv = train_model(model_conv, criterion, optimizer_conv,exp_lr_scheduler, num_epochs=10)

I wonder whether this modification is correct.