RamonYeung
/ loading_pretrained_models
Created
July 12, 2018 07:46
PyTorch Loading Pre-trained Models
This file contains 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
| # 1. Directly Load a Pre-trained Model | |
| # https://github.com/pytorch/vision/tree/master/torchvision/models | |
| import torchvision.models as models | |
| resnet50 = models.resnet50(pretrained=True) | |
| # or | |
| model = models.resnet50(pretrained=False) | |
| # Maybe you want to modify the last fc layer? | |
| resnet.fc = nn.Linear(2048, 2) |