How the pytorch freeze network in some layers, only the rest of the training?

6

Hi Spandan;

I try to replicate your code on Resnet 18. Kind of completed the code. My aim was to freeze all layers in the network except the classification layer and the layer/block preceding it. Could you please let me know your thoughts if this is right

import torch
import torchvision

model = torchvision.models.resnet18(pretrained=True)

lt=8
cntr=0

for child in model.children():
cntr+=1

if cntr < lt:
	print child
	for param in child.parameters():
		param.requires_grad = False

num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs,2)

criterion = nn.CrossEntropyLoss()

optimizer_ft = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001, momentum=0.9)

3 Likes