from torchvision.models import inception_v3
# load the pre-trained model
model = inception_v3(pretrained=True)
# Freeze the weights of the earlier layers since fine-tuning
for param in model.parameters():
param.requires_grad = False
# Replace the last fully connected layer
inp_feat_lin = model.fc.in_features
model.fc = torch.nn.Linear(inp_feat_lin, 50)
# Replace the AuxLogits layer
inp_feat_aux = model.AuxLogits.fc.in_features
model.AuxLogits.fc = torch.nn.Linear(inp_feat_aux, 50)
# Set the `requires_grad` attribute to `True` for the last layer(s) since fine-tuning
model.fc.requires_grad = True
model.AuxLogits.fc.requires_grad = True
1条答案
按热度按时间8cdiaqws1#
您可以通过将预训练版本中的两个层修改为以下内容来实现相同效果:
把这个添加到你的类模型的定义中,你可以在pytorch教程页面上深入阅读更多。