Skip to content

Instantly share code, notes, and snippets.

@dansuh17
Created February 15, 2018 05:34
Show Gist options
  • Save dansuh17/ef19051b75c363063b361c7a5462db05 to your computer and use it in GitHub Desktop.
Save dansuh17/ef19051b75c363063b361c7a5462db05 to your computer and use it in GitHub Desktop.
Finetuning pretrained pytorch model
# from : http://pytorch.org/docs/master/notes/autograd.html
# pretrained models here : https://github.com/pytorch/vision/tree/master/torchvision/models
model = torchvision.models.resnet18(pretrained=True)
for param in model.parameters():
param.requires_grad = False
# Replace the last fully-connected layer
# Parameters of newly constructed modules have requires_grad=True by default
model.fc = nn.Linear(512, 100)
# Optimize only the classifier
optimizer = optim.SGD(model.fc.parameters(), lr=1e-2, momentum=0.9)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment