Freezing and thawing of pytoch

Load part of the pre training model

In fact, most of the time, we need to adjust our model according to our tasks , Therefore, it is difficult to ensure that the model is exactly the same as the public model , But the parameters of the pre training model do help to improve the accuracy of training , In order to combine the advantages of the two , We need to load some pre training models .

pretrained_dict = model_zoo.load_url(model_urls['resnet152'])
model_dict = model.state_dict()
#  take pretrained_dict It doesn't belong to model_dict Key out of 
pretrained_dict = {
    k: v for k, v in pretrained_dict.items() if k in model_dict}
#  Update existing model_dict
model_dict.update(pretrained_dict)
#  Load what we really need state_dict
model.load_state_dict(model_dict)

Because we need to eliminate the mismatched keys in the original model , That is, the name of the layer , Therefore, the changed layer of our new model needs to be different from the name of the corresponding layer of the original model , such as ：resnet The name of the last floor is fc(PyTorch in ), So what we modified resnet You can't take this name on the last floor of , It can be called fc_

---

Micro change basic model

PyTorch Medium torchvision There are many commonly used models in , Can be called directly ：

• AlexNet
• VGG
• ResNet
• SqueezeNet
• DenseNet

import torchvision.models as models
resnet18 = models.resnet18()
alexnet = models.alexnet()
squeezenet = models.squeezenet1_0()
densenet = models.densenet_161()

But for our task, some layers are not directly usable , We need to change it a little , such as ,resnet The final full connection layer is divided into 1000 class , And we only 21 class ; And such as ,resnet The channel of convolution reception in the first layer is 3, The channel we may enter the picture is 4, Then it can be modified by the following methods ：

resnet.conv1 = nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3, bias=False)
resnet.fc = nn.Linear(2048, 21)

---

Step by step a thaw ( stay ECCV2020 The training method proposed in )

Load one FeatureExtractor

With densenet Of dense block For example :

dense = densenet121(pretrained=True)
self.loss_network = nn.Sequential(*list(dense.features)[:5])
self.loss_network.apply(add_sn)
for param in self.loss_network.parameters():
param.requires_grad = False

# vgg=vgg16(pretrained=True)
# self.loss_network = nn.Sequential(*list(vgg.features)[:16]).eval()
# self.loss_network.apply(add_sn)
# for param in self.loss_network.parameters():
# param.requires_grad = True

Set up a list, Store frozen layers , Convenient for later thawing

self.grad = []
for name, value in self.loss_network.named_parameters():
   	print('name: {0},\t grad: {1}'.format(name, value.requires_grad))
if value.requires_grad == False:
  	self.grad.append(name)

thaw ,unfreeze1 function ( Thaw a layer from bottom to top )

def unfreeze1(self):
		if len(self.grad)==0:
				print("All of network have been unfreeze!")
    		return
 		self.grad.pop()
  	print("-------------------")
  	for name, value in self.loss_network.named_parameters():
    		if name not in self.grad:
      			value.requires_grad = True
		for name, value in self.loss_network.named_parameters():
  		print('name: {0},\t grad: {1}'.format(name, value.requires_grad))
		print("---------------------")

Show freeze :

def showfreeze(self):
		for i in self.grad:
    		print(i)
		print("-------------------")

Thaw during training

network.unfreeze1()