用pytorch实现一个简单的图像分类器

 最近要做一个几十个样本量的图像分类，为了避免过拟合只能使用一个层数较少的网络，当然这个比较简单，可以供初学者作为一个入门的教学。

---

数据载入

 为了利用pytorch的torchvision模块提供的通用数据加载函数，保证图片文件夹的格式如：

├── images
    ├── train
        ├── class1
            ├── img1.jpg
            ├──      ...
        ├── class2
        ├──    ...
    ├── test
        ├── class1
        ├── class2
        ├──    ...

 然后就可以比较简单地直接利用ImageFolder函数载入图片，我将整个数据载入写成了一个ImageDataset.py文件：

'''
负责训练及测试数据的读取
'''

from torchvision import transforms, datasets
import os
import torch
from PIL import Image

def readImg(path):
    '''
    用于替代ImageFolder的默认读取图片函数，以读取单通道图片
    '''
    return Image.open(path)

def ImageDataset(args):
    # 数据增强及归一化
    # 图片都是100x100的，训练时随机裁取90x90的部分，测试时裁取中间的90x90
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomCrop(90),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5])
        ]),
        'test': transforms.Compose([
            transforms.CenterCrop(90),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5])
        ]),
    }

    data_dir = args.data_dir
    image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                        data_transforms[x], loader=readImg)
                        for x in ['train', 'test']}
    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x],
                    batch_size=args.batch_size, shuffle=(x == 'train'),
                    num_workers=args.num_workers)
                    for x in ['train', 'test']}
    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'test']}
    class_names = image_datasets['train'].classes

    return dataloaders, dataset_sizes, class_names

 从代码中也可以看到，我们加载图片的一个顺序大概是：

 —定义应用在图片上的各种transforms操作，具体到代码中的data_transforms。

 其中对于用于训练的图片是在原图上随机截取$90\times 90$大小的块（RandomCrop），然后做一个随机的水平翻转（RandomHorizontalFlip），转为Tensor类型（ToSensor，其内部是先将数据转为浮点类型后各像素点除以255），最后做一个归一化操作（Normalize，其实就是减去均值后除以标准差）。

 然后对于用于测试的图片的处理有所不同，主要还是直接截取原图中间部分的$90\times 90$大小的块（CenterCrop），取消水平翻转，其他的都一样。

 另外，因为我用到的图片都是$100\times 100$大小的，而且因为需要分类的对象用一个$90\times 90$大小的块也足够，所以才做了这样的预处理，如果你的图片大小不是这样的，需要在这里按照你自己的需求更改处理的方式，或者可以在最前面加一个resize操作将你的图片缩放为$100\times 100$大小。

 —用torchvision提供的数据加载函数加载数据集，具体到代码中的image_datasets。

 这里最主要还是ImageFolder函数的应用，需要注意的是如果图片都是单通道的，那就需要自己定义上面的readImg函数，如果你的图片是三通道图片（彩色图）那就不用了，直接在ImageFolder里面使用默认的就行了（直接删去loader=readImg）。

 —创建DataLoader的实例以为后面的训练提供数据

 这里就是为训练和测试阶段创建数据加载器，在DataLoader中指定了数据的来源（其实就是上一步生成的）、batch的大小（根据内存大小适当调整）、是否进行shuffle（只在训练阶段进行）以及进程数。

定义网络结构

 前面说过，用来训练的图片比较少，所以定义的网络层数也不能太多，这里用了三个卷积层来提取特征，最后再加一个全连接层作为分类器。整个网络的定义写成了一个SimpleNet.py文件:

'''
简单的用于分类的网络
'''

import torch

import torch.nn as nn

class SimpleNet(nn.Module):

    def __init__(self):
        super(SimpleNet, self).__init__()
        # 三个卷积层用于提取特征
        # 1 input channel image 90x90, 8 output channel image 44x44
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 8 input channel image 44x44, 16 output channel image 22x22
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 16 input channel image 22x22, 32 output channel image 10x10
        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 分类
        self.classifier = nn.Sequential(
            nn.Linear(32 * 10 * 10, 3)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = x.view(-1, 32 * 10 * 10)
        x = self.classifier(x)
        return x

 实际上上面定义的网络结构就是一个小的金字塔结构，通道数逐渐增大的同时feature map的大小在逐渐减小。每个block里面都是卷积(Conv2d)-激活函数(ReLU)-最大池化(MaxPool2d)的这种较为普遍的形式，在前向传播(forward)里也就是把这些给组合起来罢了，唯一需要注意的也就是在从最后一个卷积层到全连接层之间需要有一个view的过程，将多个通道输出的feature map给展成一个一维向量的形式作为全连接层的输入。另外，最后一个全连接层的输出数即为类别数，这个需要根据实际数据的类别数作出适当修改。

训练网络

 最后就是训练网络的部分了，整个训练部分的代码参考了pytorch官方教程中的实现方式，整合成一个train.py文件：：

'''
进行训练
'''

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import time
import os
import json
from math import ceil
import argparse
import copy
from ImageDataset import ImageDataset
from SimpleNet import SimpleNet
from tensorboardX import SummaryWriter
writer = SummaryWriter(log_dir='log')


def train_model(args, model, criterion, optimizer, scheduler, num_epochs, dataset_sizes, use_gpu):

    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    device = torch.device('cuda' if use_gpu else 'cpu')

    for epoch in range(args.start_epoch, num_epochs):
        # 每一个epoch中都有一个训练和一个验证过程(Each epoch has a training and validation phase)
        for phase in ['train', 'test']:
            if phase == 'train':
                scheduler.step(epoch)
                # 设置为训练模式(Set model to training mode)
                model.train()
            else:
                # 设置为验证模式(Set model to evaluate mode)
                model.eval()

            running_loss = 0.0
            running_corrects = 0

            tic_batch = time.time()

            # 在多个batch上依次处理数据(Iterate over data)
            for i, (inputs, labels) in enumerate(dataloders[phase]):
                inputs = inputs.to(device)
                labels = labels.to(device)

                # 梯度置零(zero the parameter gradients)
                optimizer.zero_grad()

                # 前向传播(forward)
                # 训练模式下才记录梯度以进行反向传播(track history if only in train)
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)
                    # 训练模式下进行反向传播与梯度下降(backward + optimize only if in training phase)
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # 统计损失和准确率(statistics)
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

                batch_loss = running_loss / (i * args.batch_size + inputs.size(0))
                batch_acc = running_corrects.double() / (i * args.batch_size + inputs.size(0))

                if phase == 'train' and (i + 1) % args.print_freq == 0:
                    print('[Epoch {}/{}]-[batch:{}/{}] lr:{:.6f} {} Loss: {:.6f}  Acc: {:.4f}  Time: {:.4f} sec/batch'.format(
                          epoch + 1, num_epochs, i + 1, ceil(dataset_sizes[phase]/args.batch_size), scheduler.get_lr()[0], phase, batch_loss, batch_acc, (time.time()-tic_batch)/args.print_freq))
                    tic_batch = time.time()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            if epoch == 0:
                os.remove('result.txt')
            with open('result.txt', 'a') as f:
                f.write('Epoch:{}/{} {} Loss: {:.4f} Acc: {:.4f} \n'.format(epoch + 1, num_epochs, phase, epoch_loss, epoch_acc))

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
            
            writer.add_scalar(phase + '/Loss', epoch_loss, epoch)
            writer.add_scalar(phase + '/Acc', epoch_acc, epoch)

        if (epoch + 1) % args.save_epoch_freq == 0:
            if not os.path.exists(args.save_path):
                os.makedirs(args.save_path)
            torch.save(model.state_dict(), os.path.join(args.save_path, "epoch_" + str(epoch) + ".pth"))
        
        # 深拷贝模型(deep copy the model)
        if phase == 'test' and epoch_acc > best_acc:
            best_acc = epoch_acc
            best_model_wts = copy.deepcopy(model.state_dict())

    # 将model保存为graph
    writer.add_graph(model, (inputs,))

    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Accuracy: {:4f}'.format(best_acc))

    # 载入最佳模型参数(load best model weights)
    model.load_state_dict(best_model_wts)
    return model

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='classification')
    # 图片数据的根目录(Root catalog of images)
    parser.add_argument('--data-dir', type=str, default='images')
    parser.add_argument('--batch-size', type=int, default=16)
    parser.add_argument('--num-epochs', type=int, default=150)
    parser.add_argument('--lr', type=float, default=0.045)
    parser.add_argument('--num-workers', type=int, default=4)
    parser.add_argument('--print-freq', type=int, default=1)
    parser.add_argument('--save-epoch-freq', type=int, default=1)
    parser.add_argument('--save-path', type=str, default='output')
    parser.add_argument('--resume', type=str, default='', help='For training from one checkpoint')
    parser.add_argument('--start-epoch', type=int, default=0, help='Corresponding to the epoch of resume')
    args = parser.parse_args()

    # read data
    dataloders, dataset_sizes, class_names = ImageDataset(args)

    with open('class_names.json', 'w') as f:
        json.dump(class_names, f)

    # use gpu or not
    use_gpu = torch.cuda.is_available()
    print("use_gpu:{}".format(use_gpu))

    # get model
    model = SimpleNet()

    if args.resume:
        if os.path.isfile(args.resume):
            print(("=> loading checkpoint '{}'".format(args.resume)))
            model.load_state_dict(torch.load(args.resume))
        else:
            print(("=> no checkpoint found at '{}'".format(args.resume)))

    if use_gpu:
        model = torch.nn.DataParallel(model)
        model.to(torch.device('cuda'))
    else:
        model.to(torch.device('cpu'))

    # 用交叉熵损失函数(define loss function)
    criterion = nn.CrossEntropyLoss()

    # 梯度下降(Observe that all parameters are being optimized)
    optimizer_ft = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=0.00004)

    # Decay LR by a factor of 0.98 every 1 epoch
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=1, gamma=0.98)

    model = train_model(args=args,
                        model=model,
                        criterion=criterion,
                        optimizer=optimizer_ft,
                        scheduler=exp_lr_scheduler,
                        num_epochs=args.num_epochs,
                        dataset_sizes=dataset_sizes,
                        use_gpu=use_gpu)

    torch.save(model.state_dict(), os.path.join(args.save_path, 'best_model_wts.pth'))

    writer.close()

 我们从main部分开始看起，首先是建立一个parser传递一些参数，再然后利用第一部分建立的ImageDataset载入图片数据，接着根据硬件情况决定是否使用GPU，建立模型后根据是否传入模型参数文件决定是否从之前某个epoch断点继续训练，紧接着需要为训练部分提供损失函数（这里使用交叉熵损失函数）、优化器（这里使用SGD即随机梯度下降）以及学习率调整方法（这里使用等间隔调整学习率），那有了以上的准备就可以开始训练过程了。

 训练过程主要见函数train_model()，实际上参考了pytorch官方教程。在其上增加了利用tensorboardX记录训练过程中loss和accuracy的变化情况以及可视化网络结构的代码，训练完后可以利用tensorboard --logdir log查看训练过程中loss和accuracy的变化情况以及可视化网络结构。

 下图所示是网络的结构：

网络结构

 下图所示是在我自己的数据上训练了150个epoch的结果：

训练结果

测试

 训练完之后，可以挑选表现结果最好的一个epoch的参数，重新载入以进行测试，这部分整合成一个test.py文件：

'''
测试分类
'''

from PIL import Image
from torchvision import transforms
import torch
from torch.autograd import Variable
import os
import json
from SimpleNet import SimpleNet

def predict_image(model, image_path):
    image = Image.open(image_path)

    # 测试时截取中间的90x90
    transformation1 = transforms.Compose([
        transforms.CenterCrop(90),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])

    ])

    # 预处理图像 
    image_tensor = transformation1(image).float()

    # 额外添加一个批次维度，因为PyTorch将所有的图像当做批次
    image_tensor = image_tensor.unsqueeze_(0)

    if torch.cuda.is_available():
        image_tensor.cuda()

    # 将输入变为变量
    input = Variable(image_tensor)

    # 预测图像的类别
    output = model(input)

    index = output.data.numpy().argmax()

    return index

if __name__ == '__main__':

    best_model_path = './output/epoch_462.pth'
    model = SimpleNet()
    model.load_state_dict(torch.load(best_model_path))
    model.eval()

    with open('class_names.json', 'r') as f:
        class_names = json.load(f)

    img_path = './images/test/bubble/066.jpg'
    predict_class = class_names[predict_image(model, img_path)]
    print(predict_class)

源代码

 完整的代码可见我的github项目–SimpleNet。