pytorch加载VGG16及进行fine-tuning训练

1. 加载VGG16模型并打印查看

from torchvision import models
net=models.vgg16()
print(net)

1.1结果说明

1.2查看某一部分

**

2. 加载模型进行预训练，改变classifier层，固定feature层参数

**
2.1模型搭建

import torch
import torch.nn as nn
from torchvision import models
from torchsummary import summary
net=models.vgg16()

class VGGnet(nn.Module):
    def __init__(self,feature_extract=True,num_classes=5):
        super(VGGnet, self).__init__()
        #导入VGG16模型
        model = models.vgg16(pretrained=True)
        #加载features部分
        self.features = model.features
        #固定特征提取层参数
        set_parameter_requires_grad(self.features, feature_extract)
        #加载avgpool层
        self.avgpool=model.avgpool
        #改变classifier：分类输出层
        self.classifier = nn.Sequential(
            nn.Linear(512*7*7 , 1024),
            nn.ReLU(),
            nn.Linear(1024, 1024),
            nn.ReLU(),
            nn.Linear(1024, num_classes)
        )
        
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), 512*7*7)
        out=self.classifier(x)
        return out
        
 #固定参数，不进行训练   
def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False

2.2 模型对比

net_self=VGGnet()
    print('model_build','**'*20)
    print(net_self)

结果：与原始模型（1.1中的结果）改变的是classifier层

3. 模型训练
   3.1数据及数据加载代码来源：Tensorflow2.1.0 自定义数据集：精灵宝可梦数据集

import os
from torch.utils.data import Dataset, DataLoader  #自定义的母类，必须的
from torchvision.transforms import transforms
from PIL import Image 
import torch
import glob
import csv
import random

class Pokemon(Dataset):
    def __init__(self, root, resize, mode):
        super(Pokemon, self).__init__()
        self.root = root
        self.resize = resize
        self.name2label = {}  # "sq...":0
        for name in sorted(os.listdir(os.path.join(root))):
            if not os.path.isdir(os.path.join(root, name)):
                continue
            self.name2label[name] = len(self.name2label.keys()) #将英文标签名转化数字0-4
        # print(self.name2label)
        # image, label
        self.images, self.labels = self.load_csv('images.csv')  #csv文件存在 直接读取
        if mode == 'train':  # 60%                   
            self.images = self.images[:int(0.6 * len(self.images))]
            self.labels = self.labels[:int(0.6 * len(self.labels))]
        elif mode == 'val':  # 20% = 60%->80%
            self.images = self.images[int(
                0.6 * len(self.images)):int(0.8 * len(self.images))]
            self.labels = self.labels[int(
                0.6 * len(self.labels)):int(0.8 * len(self.labels))]
        else:  # 20% = 80%->100%
            self.images = self.images[int(0.8 * len(self.images)):]
            self.labels = self.labels[int(0.8 * len(self.labels)):]

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        # idx~[0~len(images)]
        # self.images, self.labels
        # img: 'pokemon\\bulbasaur\\00000000.png'
        # label: 0
        img, label = self.images[idx], self.labels[idx]
        
        tf = transforms.Compose([   #常用的数据变换器
					            lambda x:Image.open(x).convert('RGB'),  # string path= > image data 
					            #这里开始读取了数据的内容了
					            transforms.Resize(   #数据预处理部分
					                (int(self.resize * 1.25), int(self.resize * 1.25))), 
					            transforms.RandomRotation(15), 
					            transforms.CenterCrop(self.resize), #防止旋转后边界出现黑框部分
					            transforms.ToTensor(),
					            transforms.Normalize(mean=[0.485, 0.456, 0.406],
					                                 std=[0.229, 0.224, 0.225])
       							 ])
        img = tf(img)
        label = torch.tensor(label)  #转化tensor
        return img, label       #返回当前的数据内容和标签
    
    def load_csv(self, filename):
        if not os.path.exists(os.path.join(self.root, filename)): 
        	     #如果没有保存csv文件，那么我们需要写一个csv文件，如果有了直接读取csv文件
            images = []
            for name in self.name2label.keys():   
            	            # 'pokemon\\mewtwo\\00001.png
                images += glob.glob(os.path.join(self.root, name, '*.png'))
                images += glob.glob(os.path.join(self.root, name, '*.jpg'))
                images += glob.glob(os.path.join(self.root, name, '*.jpeg'))
                	
            	        
            random.shuffle(images)
            with open(os.path.join(self.root, filename), mode='w', newline='') as f:
                writer = csv.writer(f)
                for img in images:  # 'pokemon\\bulbasaur\\00000000.png'
                    name = img.split(os.sep)[-2]        #从名字就可以读取标签
                    label = self.name2label[name]
                    	                # 'pokemon\\bulbasaur\\00000000.png', 0
                    writer.writerow([img, label])  #写进csv文件
        	
        	    # read from csv file
        images, labels = [], []
        with open(os.path.join(self.root, filename)) as f:
            reader = csv.reader(f)
            for row in reader:
            	            # 'pokemon\\bulbasaur\\00000000.png', 0
                img, label = row
                label = int(label)
                images.append(img)
                labels.append(label)
        assert len(images) == len(labels)
        return images, labels
    
    def denormalize(self, x_hat):
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        # x_hat = (x-mean)/std
        # x = x_hat*std = mean
        # x: [c, h, w]
        # mean: [3] => [3, 1, 1]
        mean = torch.tensor(mean).unsqueeze(1).unsqueeze(1)
        std = torch.tensor(std).unsqueeze(1).unsqueeze(1)
        # print(mean.shape, std.shape)
        x = x_hat * std + mean
        return x
if __name__=='__main__':
    db = Pokemon('pokeman', 224, 'train')
    loader = DataLoader(db, batch_size=32, shuffle=True)
    for x, y in loader: #此时x,y是批量的数据
        print(x.shape)

3.2 训练

import torch
import torch.nn as nn
from torchvision import models

class VGGnet(nn.Module):
    def __init__(self,feature_extract=True,num_classes=5):
        super(VGGnet, self).__init__()
        model = models.vgg16(pretrained=True)
        self.features = model.features
        set_parameter_requires_grad(self.features, feature_extract)#固定特征提取层参数
        self.avgpool=model.avgpool
        self.classifier = nn.Sequential(
            nn.Linear(512*7*7 , 1024),
            nn.ReLU(),
            nn.Linear(1024, 1024),
            nn.ReLU(),
            nn.Linear(1024, num_classes)
        )
        
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), 512*7*7)
        out=self.classifier(x)
        return out
    
def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False

if __name__=="__main__":
    import torch.nn as nn
    from torch.utils.data import DataLoader
    from data_read import Pokemon
    
    # In[]
    learning_rate=0.001
    num_epochs = 2               # train the training data n times, to save time, we just train 1 epoch
    batch_size = 32
    LR = 0.01              # learning rate
    # In[]
    
    train_dataset = Pokemon('pokeman', 224, 'train')
    val_dataset = Pokemon('pokeman', 224, 'val')
    test_dataset = Pokemon('pokeman', 224, 'test')
    
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    test_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
    val_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
        
    
    # In[]
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model=VGGnet().to(device)
     
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    
    # In[]
    total_step = len(train_loader)
    for epoch in range(num_epochs):
        for i, (images, labels) in enumerate(train_loader):
    
            images = images.to(device)
            labels = labels.to(device)
     
            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)
     
            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
     
            if (i + 1) % 2 == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                      .format(epoch + 1, num_epochs, i + 1, total_step, loss.item()))
    
    # In[]      
    # Test the model
    model.eval()  #
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
        print('Test Accuracy  {} %'.format(100 * correct / total))

3.3 结果

4. 数据集链接：所有分享的数据集都在这个文件夹，分类，迁移学习，图像分割等
   数据集名称：VGG16(fine-tuning_pokeman)
   百度云链接：https://pan.baidu.com/s/1eyDFV6YVaOHwr9QrRcKXFg
   提取码：rlne