迁移学习resnet

import torch
import numpy as np
import torchvision
import torch.nn as nn

from torchvision import datasets,transforms,models
import matplotlib.pyplot as plt
import time
import os
import copy
print("Torchvision Version:",torchvision.__version__)

data_dir="./hymenoptera_data"
batch_size=32
input_size=224
model_name="resnet"
num_classes=2
num_epochs=15
feature_extract=True
data_transforms={
    "train":transforms.Compose([
        transforms.RandomResizedCrop(input_size),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.482,0.456,0.406],[0.229,0.224,0.225])
    ]),
    "val":transforms.Compose([

    transforms.RandomResizedCrop(input_size),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.482, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
image_datasets={x:datasets.ImageFolder(os.path.join(data_dir,x),data_transforms[x])
                for x in ["train",'val']}
dataloader_dict={x:torch.utils.data.DataLoader(image_datasets[x],batch_size=batch_size,
                shuffle=True)for x in ['train','val']}
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
inputs,labels=next(iter(dataloader_dict["train"]))
#print(inputs.shape)#一个batch
#print(labels)


#加载resent模型并修改全连接层
def set_parameter_requires_grad(model,feature_extract):
    if feature_extract:
        for param in model.parameters():
            param.requires_grad=False

def initialize_model(model_name,num_classes,feature_extract,use_pretrained=True):
    if model_name=="resnet":
        model_ft=models.resnet18(pretrained=use_pretrained)
        set_parameter_requires_grad(model_ft,feature_extract)
        num_ftrs=model_ft.fc.in_features
        model_ft.fc=nn.Linear(num_ftrs,num_classes)
        input_size=224
    else:
        print("model not implemented")
        return None,None

    return model_ft,input_size
model_ft,input_size=initialize_model(model_name,num_classes,feature_extract,use_pretrained=True)
#print(model_ft)
print('-'*200)


def train_model(model,dataloaders,loss_fn,optimizer,num_epochs):
    best_model_wts=copy.deepcopy(model.state_dict)
    best_acc=0.
    val_acc_history=[]
    for epoch in range(num_epochs):
        for phase in ["train","val"]:
            running_loss=0.
            running_corrects=0.
            if phase=="train":
                model.train()
            else:
                model.eval()

            for inputs,labels in dataloaders[phase]:
                inputs,labels=inputs.to(device),labels.to(device)

                with torch.autograd.set_grad_enabled(phase=="train"):
                    outputs=model(inputs)
                    loss=loss_fn(outputs,labels)
                preds=outputs.argmax(dim=1)
                if phase=="train":
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
                running_loss+=loss.item()*inputs.size(0)
                running_corrects+=torch.sum(preds.view(-1)==labels.view(-1)).item()

            epoch_loss=running_loss/len(dataloaders[phase].dataset)
            epoch_acc=running_corrects/len(dataloaders[phase].dataset)

            print("Phase{}  loss:{},   acc:{}".format(phase,epoch_loss,epoch_acc))

            if phase=="val" and epoch_acc>best_acc:
                best_acc=epoch_acc
                best_model_wts=copy.deepcopy(model.state_dict())
            if phase=="val":
                val_acc_history.append(epoch_acc)
    model.load_state_dict(best_model_wts)
    return  model,val_acc_history

model_ft=model_ft.to(device)
optimizer=torch.optim.SGD(filter(lambda  p: p.requires_grad,model_ft.parameters()),
                          lr=0.001,momentum=0.9)
loss_fn=nn.CrossEntropyLoss()
print("feature extraction: 我们不再改变训练模型的参数，而是只更新我们改变过的部分模型参数。"
      "我们之所以叫它feature extraction是因为我们把预训练的CNN模型当做一个特征提取模型，利用提取出来的特征做来完成我们的训练任务。")
_,ohist=train_model(model_ft,dataloader_dict,loss_fn,optimizer,num_epochs=num_epochs)

print("-"*200)


model_scratch,_=initialize_model(model_name,num_classes,feature_extract=False,use_pretrained=False)
model_scratch=model_ft.to(device)
optimizer=torch.optim.SGD(filter(lambda  p: p.requires_grad,model_ft.parameters()),
                          lr=0.001,momentum=0.9)
loss_fn=nn.CrossEntropyLoss()
print("fine tuning: 从一个预训练模型开始，我们改变一些模型的架构，然后继续训练整个模型的参数。")
_,scratch_ohist=train_model(model_ft,dataloader_dict,loss_fn,optimizer,num_epochs=num_epochs)

plt.title("Accuracy vs. Training Epoch")
plt.xlabel("Training Epoch")
plt.ylabel("Accuracy")
plt.plot(range(1,num_epochs+1),ohist,label="Pretrained")
plt.plot(range(1,num_epochs+1),scratch_ohist,label="No_pretrained")
plt.ylim((0,1.))
plt.xticks(np.arange(1,num_epochs+1,1.0))
plt.legend()
plt.show()