resnet pre-training model and fine-tuning

Fine-tuning:
Jointly training the newly added classifier and some or all convolutional layers allows us to fine-tune the high-order feature representations in the base model as they are relevant to the specific task.
Only when the classifier is trained (that is, the Linear layer) can the convolution base be fine-tuned. Otherwise, the training error will be very large at the beginning, and what was learned in the convolution layer before fine-tuning will be destroyed.
 
Steps:
1. Add a custom layer to the pre-trained convolution base
2. Freeze the convolution base
3. Train the added classification layer
4. Unfreeze part of the convolution base or all of it (generally unfreeze the convolution base close to the output part)
5. Jointly train the entire model

General steps for transfer learning :

 Code:

"""Use Resnet model to pre-train the model + fine-tuning""" 


import torch.cuda 
import torchvision 
from torch import nn 
import os 
from torchvision import transforms 
# import matplotlib.pyplot as plt 


#Preprocess data 
base_dir = r'./dataset/4weather ' 
train_dir = os.path.join(base_dir,'train') 
test_dir = os.path.join(base_dir,'test') 

train_transformer = torchvision.transforms.Compose( 
     [ transforms.Resize((224,224)), 
     transforms.RandomCrop ((192,192)), 
     transforms.ColorJitter(brightness=0.4), 
     transforms.RandomHorizontalFlip(0.2), 
     transforms.ColorJitter(contrast=4), 
     transforms.ToTensor(), 
     transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5,0.5, 0.5])
     ]

)

test_transformer = torchvision.transforms.Compose([
                         transforms.Resize((192,192)),
                          transforms.ToTensor(),
                          transforms.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5])])

#创建dataset
train_ds = torchvision.datasets.ImageFolder(
    train_dir,
    transform = train_transformer
)

test_ds = torchvision.datasets.ImageFolder(
    test_dir,
    transform =test_transformer
)

#创建dataloader
batch = 8
train_dl =torch .utils.data.DataLoader(
    train_ds, batch_size = batch, 
    shuffle = True 
) 
test_dl = torch.utils.data.DataLoader( 
    test_ds, 
    batch_size = batch)

imgs,labels =next(iter(train_dl)) 



#Load model 
device = 'cuda' if torch.cuda.is_available() else 'cpu' 
model = torchvision.models.resnet18().to(device) 


for param in model. parameters(): 
    param.requires_grad = False 

#Only the number of classification categories of the full link layer needs to be changed 
in_f= model.fc.in_features 
#Replacing the full link layer is trainable 
model.fc = nn.Linear(in_f, 4) 

#The optimizer only needs to optimize the last layer 
optim =torch .optim.Adam(model.fc.parameters(),lr =0.001) 
#Loss function 
loss_fn = nn.CrossEntropyLoss() 

#The training function fit must specify model.train, There is a BN layer in model.eval Resnet 

def fit(epoch,model,trainloader,testloader): 
    correct = 0 
    total = 0 
    running_loss =0 
    model.train() #Indicate that this is the train mode and requires bn and drop
    for x,y in trainloader:
        if torch.cuda.is_available():
            x,y =x.to('cuda'),y.to('cuda')
        y_pred =model(x)
        loss = loss_fn(y_pred,y)
        optim.zero_grad()
        loss.backward()
        optim.step()
        with torch.no_grad():
            y_pred = torch.argmax(y_pred,dim=1)
            correct +=(y_pred==y).sum().item()
            total += y.size(0)
            running_loss += loss.item()
    epoch_loss = running_loss/len(trainloader.dataset)
    epoch_acc =correct/total

    test_correct = 0
    test_total = 0
    test_running_loss =0

    model.eval()
    with torch.no_grad():
        for x,y in testloader:
            if torch.cuda.is_available():
                x,y = x.to('cuda'),y.to('cuda')
            y_pred =model(x)
            loss = loss_fn(y_pred,y)
            y_pred = torch.argmax(y_pred,dim=1)
            test_correct +=(y_pred==y).sum().item()
            test_total +=y.size(0)
            test_running_loss +=loss.item()
    epoch_tst_loss =test_running_loss/len(testloader.dataset)
    epoch_tst_acc = test_correct/test_total

    print('epoch',epoch,'loss',round(epoch_loss,3),
          'acc:',round(epoch_acc,3),
          'test_loss:',round(epoch_tst_loss,3),
          'test_acc:',round(epoch_tst_acc,3))
    return epoch_loss ,epoch_acc,epoch_tst_loss,epoch_tst_acc

# epochs =5 
# train_loss =[] 
# train_acc =[] 
# test_loss =[] 
# test_acc=[] 
# 
# for epoch in range(epochs): 
# epoch_loss,epoch_acc,epoch_tst_loss,epoch_tst_acc =fit(epoch,model,train_dl ,test_dl) 
# train_loss.append(epoch_loss) 
# train_acc.append(epoch_acc) 
# test_loss.append(epoch_tst_loss) 
# test_acc.append(epoch_tst_acc)


#绘制图像：
# plt.plot(range(1,epochs+1),train_loss,label='train_loss')
# plt.plot(range(1,epochs+1),test_loss,label = 'test_loss' )
# plt.legend() 
# plt.show() 
# #Fine-tuning 
for param in model.parameters(): 
    param.requires_grad=True 

extend_epoch =8 
#The learning rate should be smaller during fine-tuning 
optimizer = torch.optim.Adam( model.parameters(),lr=0.00001) 
# 
# #Training process 
train_loss =[] 
train_acc =[] 
test_loss =[] 
test_acc=[] 

for epoch in range(extend_epoch): 
    epoch_loss,epoch_acc,epoch_tst_loss,epoch_tst_acc =fit(epoch ,model,train_dl,test_dl) 
    train_loss.append(epoch_loss) 
    train_acc.append(epoch_acc) 
    test_loss.append(epoch_tst_loss) 
    test_acc.append(epoch_tst_acc)