基于ModelArts平台利用VGG网络实现CIFAR10图像分类任务

CIFAR-10 是由 Hinton 的学生 Alex Krizhevsky 和 Ilya Sutskever 整理的一个用于识别普适物体的小型数据集。

一共包含 10 个类别的 RGB 彩色图 片：飞机（ airlane ）、汽车（ automobile ）、鸟类（ bird ）、猫（ cat ）、鹿（ deer ）、狗（ dog ）、蛙类（ frog ）、马（ horse ）、船（ ship ）和卡车（ truck ）。图片的尺寸为 32×32 ，数据集中一共有 50000 张训练圄片和 10000 张测试图片。

本次实验在华为云ModelArts平台-开发环境-NoteBook完成。

下载和缓存数据集

这里实现了几个函数来下载 CIFAR-10 数据集， CIFAR_URL是下载数据集的网址。

import os
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms
import torch.optim as optim
import time
import torchvision
CIFAR_URL = 'https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz'

下面的download函数用来下载并解压缩数据集， 将数据集缓存在本地目录（默认情况下为…/data）中。 如果缓存目录中已经存在此数据集文件，并且其名称与下载文件名匹配， 我们将使用缓存的文件，以避免重复的下载。

下载完成后函数执行解压缩，为避免重复解压缩，同样检查是否有文件名相匹配。

def downloadCIFAR():
    if not os.path.exists('../data'):
        os.system('mkdir ../data')
        print('make dic')
    if not os.path.exists('../data/cifar-10-python.tar.gz'):
        print('downloading CIFAR')
        os.system('wget {}'.format(CIFAR_URL))
        os.system('mv ./cifar-10-python.tar.gz ../data/cifar-10-python.tar.gz')
    else:
        print('CIFAR exist')
    if not os.path.exists('../data/cifar-10-batches-py'):
        print('unpacking CIFAR')
        os.system('tar -zxvf ../data/cifar-10-python.tar.gz -C ../data/')
    else:
        print('CIFAR unpacked')
        
downloadCIFAR()

数据处理

下面的代码用作处理数据。

CIFAR数据保存为 pkl 格式，需要加载 pickle 来读取文件，下面的 unpickle 函数实现了读取 pkl 文件并返回 List 的功能。

可以尝试展示一下读取的数据，下面的 show_img 函数实现了读取数据并展示图片和label的功能

transform_train = transforms.Compose(
    [transforms.Pad(4),
     transforms.ToTensor(),
     transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
     transforms.RandomHorizontalFlip(),
     transforms.RandomGrayscale(),
     transforms.RandomCrop(32, padding=4),
])

def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict

def show_img(num):
    dict = unpickle('../data/cifar-10-batches-py/data_batch_1')
    info = unpickle('../data/cifar-10-batches-py/batches.meta')
    np_img = dict[b'data'][num].reshape([3, 32, 32])
    np_img = np_img.transpose(1, 2, 0)
    img=Image.fromarray(np_img)
    img=transform_train(img)
    img=img.swapaxes(0,1)
    img=img.swapaxes(2,1)
    np_img=np.array(img)
    label = info[b'label_names'][dict[b'labels'][num]]
    
    plt.imshow(np_img)
    plt.axis('on')
    plt.title(label)
    plt.show()

    
show_img(6)

下面的代码实现了将读取到的数据转为torch.tensor以及一系列形状操作。

transform_train = transforms.Compose(
    [transforms.Pad(4),
     transforms.ToTensor(),
     transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
     transforms.RandomHorizontalFlip(),
     transforms.RandomGrayscale(),
     transforms.RandomCrop(32, padding=4),
])

transform_test = transforms.Compose(
    [
     transforms.ToTensor(),
     transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]
)

def get_label_data(train_val, num=0):
    if train_val == 'train':
        dict = unpickle('../data/cifar-10-batches-py/data_batch_{}'.format(num))
        
        # for Num in range(len(dict[b'data'])):
        #     np_img = dict[b'data'][Num].reshape([3, 32, 32])
        #     np_img = np_img.transpose(1, 2, 0)
        #     img=Image.fromarray(np_img)
        #     img=transform_train(img)
        #     np_img=np.array(img)
        #     img=img.swapaxes(0,1)
        #     img=img.swapaxes(2,1)
        #     dict[b'data'][Num]=np_img.reshape(3072)
            

        return torch.tensor(dict[b'labels'], dtype=torch.long), torch.tensor(dict[b'data'], dtype=torch.float32).reshape(10000, 3, 32, 32)
    else:
        dict = unpickle('../data/cifar-10-batches-py/test_batch')
        return torch.tensor(dict[b'labels'], dtype=torch.long), torch.tensor(dict[b'data'], dtype=torch.float32).reshape(10000, 3, 32, 32)

def load_data():
    for i in range(5):
        if i == 0:
            train_labels, train_data = get_label_data('train', i + 1)
        else:
            res = get_label_data('train', i + 1)
            train_labels = torch.cat((train_labels, res[0]), dim=0)
            train_data = torch.cat((train_data, res[1]), dim=0)
    test_labels, test_data = get_label_data('test')
    return train_labels, train_data, test_labels, test_data

train_labels, train_data, test_labels, test_data = load_data()
print(train_labels.shape, train_data.shape, test_labels.shape, test_data.shape)

训练

设置trainloader与testloader

train_set = torch.utils.data.TensorDataset(train_data, train_labels)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=128, shuffle=True)
test_set = torch.utils.data.TensorDataset(test_data, test_labels)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=128, shuffle=False)

下面手动实现VGG16网络

vgg = [96, 96, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']
class VGG(nn.Module):
    def __init__(self, vgg):
        super(VGG, self).__init__()
        self.features = self._make_layers(vgg)
        self.dense = nn.Sequential(
            nn.Linear(512, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(0.4),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(0.4),
        )
        self.classifier = nn.Linear(4096, 10)
 
    def forward(self, x):
        out = self.features(x)
        out = out.view(out.size(0), -1)
        out = self.dense(out)
        out = self.classifier(out)
        return out
 
    def _make_layers(self, vgg):
        layers = []
        in_channels = 3
        for x in vgg:
            if x == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
                           nn.BatchNorm2d(x),
                           nn.ReLU(inplace=True)]
                in_channels = x
 
        layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
        return nn.Sequential(*layers)

训练时会用的各种配置，其中scheduler对最后结果的收敛起到至关重要的作用。

model = VGG(vgg)                                                                        #模型
criterion = nn.CrossEntropyLoss()                                                       #损失函数
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)     #优化器
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.4, last_epoch=-1) #学习率调整器
# 转移至GPU
device = torch.device("cuda:0")
model.to(device)

下面开始训练

train_acc_all = []
train_loss_all=[]
test_acc_all = []
test_loss_all=[]
def train(model, optimizer, num_epochs=40, device=device):
    for epoch in range(num_epochs):
        # 训练
        model.train()
        train_loss = 0.0
        train_acc = 0.0
        total = 0
        correct = 0
        for data, label in train_loader:
            data = data.to(device)
            label = label.to(device)  # 把数据放到GPU上面训练
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, label)
            loss.backward()
            optimizer.step()
            train_loss += loss.item() * data.size(0)
            _, predicted = torch.max(output.data, 1)
            total += label.size(0)
            correct += (predicted == label).sum().item()
        scheduler.step()

        train_loss = train_loss / len(train_loader.dataset)
        train_loss_all.append(train_loss)
        train_acc = correct / total
        train_acc_all.append(train_acc)

        model.eval()
        test_loss = 0.0
        test_acc = 0.0
        total = 0
        correct = 0
        with torch.no_grad():
            for data, label in test_loader:
                data = data.to(device)
                label = label.to(device)
                output = model(data)
                loss = criterion(output, label)
                test_loss += loss.item() * data.size(0)
                _, predicted = torch.max(output.data, 1)
                total += label.size(0)
                correct += (predicted == label).sum().item()

            test_loss = test_loss / len(test_loader.dataset)
            test_acc = correct / total
            test_loss_all.append(test_loss)
            test_acc_all.append(test_acc)
            
        print('Epoch [%d/%d], Train Loss: %.4f, Train Acc: %.4f, Test Loss: %.4f, Test Acc: %.4f'
            % (epoch + 1, num_epochs, train_loss, train_acc, test_loss, test_acc))

train(model, optimizer, num_epochs=40, device=device)

画出训练过程中的准确率与损失变化曲线

# 画图
import matplotlib.pyplot as plt

def plot_accuracy(train_acc_all, test_acc_all,i):
    plt.plot(train_acc_all, label='Training')
    plt.plot(test_acc_all, label='Test')
    plt.plot([0,40], [0.85, 0.85], c='b', linestyle='--')
    plt.title("Accuracy Curve ("+i+")")
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.show()
  
plot_accuracy(train_acc_all, test_acc_all,"vgg16")

def plot_loss(train_loss_all, test_loss_all,i):
    plt.plot(train_loss_all, label='Training')
    plt.plot(test_loss_all, label='Test')
    plt.title("Loss Curve ("+i+")")
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.show()
  
plot_loss(train_loss_all, test_loss_all,"vgg16")

最后写一些技术总结： 1.要让最后的结果收敛，就需要添加一个scheduler对学习率进行一个动态调整，比如每5轮就让学习率除以10之类的。 2.图像增广技术是十分有效的，可以显著提升最后的结果。 3.算力真的是对深度学习来说十分重要的一个东西，就这么简单的一个小任务，用自己的GPU跑一轮都要快一个小时。十分感谢华为云让我的训练加速。

总的来说，对比网络上可以查阅到的各种解决CIFAR10的方案，达到0.877的准确率已经是相当好的一个结果了，许多的参考方案最后的结果都只有0.7到0.8之间。一般的使用VGG网络来解决这项任务可能一般结果在0.8周围浮动，在注意了一些细节之后才达到了现在的效果。当然，如果使用resnet的话或许会有更好的结果，但是训练时间也一定是更长的。