Created
July 17, 2019 17:42
-
-
Save Hanrui-Wang/9884f45fc74a17ac6ecb99a54a61394f to your computer and use it in GitHub Desktop.
how to load a model, change the final layer and finetune all the layers?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| model_ft = models.resnet18(pretrained=True) | |
| num_ftrs = model_ft.fc.in_features | |
| model_ft.fc = nn.Linear(num_ftrs, 2) | |
| model_ft = model_ft.to(device) | |
| criterion = nn.CrossEntropyLoss() | |
| # Observe that all parameters are being optimized | |
| optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9) | |
| # Decay LR by a factor of 0.1 every 7 epochs | |
| exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1) | |
| model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, | |
| num_epochs=25) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment