[关闭]
@fangyang970206 2018-12-01T10:34:25.000000Z 字数 10724 阅读 1099

基于交通灯数据集的端到端分类

pytorch Classification Traffic_Light_Dataset

抓住11月的尾巴，这里写上昨天做的一个DL的作业吧，作业很简单，基于交通灯的图像分类，但这确是让你从0构建深度学习系统的好例子，很多已有的数据集都封装好了，直接调用，这篇文章将以pytorch这个深度学习框架一步步搭建分类系统。

软件包要求：
pytorch：0.4.0
torchsummary：pip install torchsummary
cv2: pip install opencv-python
matplotlib
numpy

所有代码托管到github上，链接如下：https://github.com/FangYang970206/TL_Dataset_Classification，git clone https://github.com/FangYang970206/TL_Dataset_Classification到本地。

1.数据集简介

数据集有10个类别，分别是红灯的圆球，向左，向右，向上和负例以及绿灯的圆球，向左，向右，向上和负例，如下图所示：
image_1cthr3cpb93e1ia7dnl1kl51tjh9.png-227.7kB
数据集的可通过如下链接进行下载：onedrive，baiduyun，google。
下完数据集后，解压到文件夹TL_Dataset_Classification-master中，得到一个新的文件夹TL_Dataset，可以看到TL_Dataset有以下目录：
image_1cthuhco8l4f1n05hen27a1ou216.png-14.6kB

2.代码实战

代码是在vscode上编写的，支持flask8，总共有9个文件，下面一一介绍。建议在看代码的时候从main.py文件开始看，大致脉络就清楚了。

2.1 model.py

对于一个深度学习系统来说，model应该是最初的想法，我们想构造什么样的模型来拟合数据集，所以先写model，代码如下：

  1. import torch.nn as nn
  2. from torchsummary import summary
  5. class A2NN(nn.Module):
  6.     def __init__(self, ):
  7.         super(A2NN, self).__init__()
  8.         self.main = nn.Sequential(
  9.             nn.Conv2d(3, 16, 3, 1, 1),
  10.             nn.BatchNorm2d(16),
  11.             nn.ReLU(inplace=True),
  12.             nn.Conv2d(16, 32, 3, 1, 1),
  13.             nn.MaxPool2d(2, 2),
  14.             nn.BatchNorm2d(32),
  15.             nn.ReLU(inplace=True),
  16.             nn.Conv2d(32, 32, 3, 1, 1),
  17.             nn.MaxPool2d(2, 2),
  18.             nn.BatchNorm2d(32),
  19.             nn.ReLU(inplace=True),
  20.             nn.Conv2d(32, 64, 3, 1, 1),
  21.             nn.MaxPool2d(2, 2),
  22.             nn.BatchNorm2d(64),
  23.             nn.ReLU(inplace=True),
  24.         )
  25.         self.linear = nn.Linear(4*4*64, 9)
  27.     def forward(self, inp):
  28.         x = self.main(inp)
  29.         x = x.view(x.shape[0], -1)
  30.         x = self.linear(x)
  31.         return x
  34. if __name__ == "__main__":
  35.     nn = A2NN()
  36.     summary(nn, (3, 32, 32))

model代码不复杂，很简单，这里不多介绍，缺少基础的朋友还请自行补基础。

2.2 dataset.py

第二步我们要构建数据集类，pytorch封装了一个torch.utils.data.Dataset的类，我们可以重载__len__和__getitem__方法，来得到自己的数据集管道，__len__方法是返回数据集的长度，__getitem__是支持从0到len(self)互斥范围内的整数索引，返回的是索引对应的数据和标签。代码如下：

  1. import torch
  2. import cv2
  3. import torch.utils.data as data
  6. class_light = {
  7.                 'Red Circle': 0,
  8.                 'Green Circle': 1,
  9.                 'Red Left': 2,
  10.                 'Green Left': 3,
  11.                 'Red Up': 4,
  12.                 'Green Up': 5,
  13.                 'Red Right': 6,
  14.                 'Green Right': 7,
  15.                 'Red Negative': 8,
  16.                 'Green Negative': 8
  17. }
  20. class Traffic_Light(data.Dataset):
  21.     def __init__(self, dataset_names, img_resize_shape):
  22.         super(Traffic_Light, self).__init__()
  23.         self.dataset_names = dataset_names
  24.         self.img_resize_shape = img_resize_shape
  26.     def __getitem__(self, ind):
  27.         img = cv2.imread(self.dataset_names[ind])
  28.         img = cv2.resize(img, self.img_resize_shape)
  29.         img = img.transpose(2, 0, 1)-127.5/127.5
  30.         for key in class_light.keys():
  31.             if key in self.dataset_names[ind]:
  32.                 label = class_light[key]
  33.         # pylint: disable=E1101,E1102
  34.         return torch.from_numpy(img), torch.tensor(label)
  35.         # pylint: disable=E1101,E1102
  37.     def __len__(self):
  38.         return len(self.dataset_names)
  41. if __name__ == '__main__':
  42.     from torch.utils.data import DataLoader
  43.     from glob import glob
  44.     import os
  46.     path = 'TL_Dataset/Green Up/'
  47.     names = glob(os.path.join(path, '*.png'))
  48.     dataset = Traffic_Light(names, (32, 32))
  49.     dataload = DataLoader(dataset, batch_size=1)
  50.     for ind, (inp, label) in enumerate(dataload):
  51.         print("{}-inp_size:{}-label_size:{}".format(ind, inp.numpy().shape,
  52.                                                     label.numpy().shape))

2.3 util.py

在上面的dataset.py中，class初始化时，传入了dataset_names，所以utils.py文件中就通过get_train_val_names函数得到训练数据集和验证数据集的names，还有一个函数是检查文件夹是否存在，不存在建立文件夹。代码如下：

  1. import os
  2. from glob import glob
  5. def get_train_val_names(dataset_path, remove_names, radio=0.3):
  6.     train_names = []
  7.     val_names = []
  8.     dataset_paths = os.listdir(dataset_path)
  9.     for n in remove_names:
  10.         dataset_paths.remove(n)
  11.     for path in dataset_paths:
  12.         sub_dataset_path = os.path.join(dataset_path, path)
  13.         sub_dataset_names = glob(os.path.join(sub_dataset_path, '*.png'))
  14.         sub_dataset_len = len(sub_dataset_names)
  15.         val_names.extend(sub_dataset_names[:int(radio*sub_dataset_len)])
  16.         train_names.extend(sub_dataset_names[int(radio*sub_dataset_len):])
  17.     return {'train': train_names, 'val': val_names}
  20. def check_folder(path):
  21.     if not os.path.exists(path):
  22.         os.mkdir(path)

2.4 trainer.py

model构造好了，数据集也准备好了，现在就需要准备如果训练了，这就是trainer.py文件的作用，trainer.py构建了Trainer类，通过传入训练的一系列参数，调用Trainer.train函数进行训练，并返回loss，代码如下：

  1. import torch.nn as nn
  2. from torch.optim import Adam
  5. class Trainer:
  6.     def __init__(self, model, dataload, epoch, lr, device):
  7.         self.model = model
  8.         self.dataload = dataload
  9.         self.epoch = epoch
  10.         self.lr = lr
  11.         self.device = device
  12.         self.optimizer = Adam(self.model.parameters(), lr=self.lr)
  13.         self.criterion = nn.CrossEntropyLoss().to(self.device)
  15.     def __epoch(self, epoch):
  16.         self.model.train()
  17.         loss_sum = 0
  18.         for ind, (inp, label) in enumerate(self.dataload):
  19.             inp = inp.float().to(self.device)
  20.             label = label.long().to(self.device)
  21.             self.optimizer.zero_grad()
  22.             out = self.model.forward(inp)
  23.             loss = self.criterion(out, label)
  24.             loss.backward()
  25.             loss_sum += loss.item()
  26.             self.optimizer.step()
  27.             print('epoch{}_step{}_train_loss_: {}'.format(epoch,
  28.                                                           ind,
  29.                                                           loss.item()))
  30.         return loss_sum/(ind+1)
  32.     def train(self):
  33.         train_loss = self.__epoch(self.epoch)
  34.         return train_loss

2.5 validator.py

trainer.py文件是用来进行训练数据集的，训练过程中，我们是需要有验证集来判断我们模型的训练效果，所以这里有validator.py文件，里面封装了Validator类，与Trainer.py类似，但不同的是，我们不训练，不更新参数，model处于eval模式，代码上会有一些跟Trainer不一样,通过调用Validator.eval函数返回loss，代码如下：

  1. import torch.nn as nn
  4. class Validator:
  5.     def __init__(self, model, dataload, epoch, device, batch_size):
  6.         self.model = model
  7.         self.dataload = dataload
  8.         self.epoch = epoch
  9.         self.device = device
  10.         self.batch_size = batch_size
  11.         self.criterion = nn.CrossEntropyLoss().to(self.device)
  13.     def __epoch(self, epoch):
  14.         self.model.eval()
  15.         loss_sum = 0
  16.         for ind, (inp, label) in enumerate(self.dataload):
  17.             inp = inp.float().to(self.device)
  18.             label = label.long().to(self.device)
  19.             out = self.model.forward(inp)
  20.             loss = self.criterion(out, label)
  21.             loss_sum += loss.item()
  22.         return {'val_loss': loss_sum/(ind+1)}
  24.     def eval(self):
  25.         val_loss = self.__epoch(self.epoch)
  26.         return val_loss

2.6 logger.py

我们想看整个学习的过程，可以通过看学习曲线来进行观察。所以这里写了一个logger.py文件，用来对训练loss和验证loss进行统计并画图。代码如下：

  1. import matplotlib.pyplot as plt
  2. import os
  5. class Logger:
  6.     def __init__(self, save_path):
  7.         self.save_path = save_path
  9.     def update(self, Kwarg):
  10.         self.__plot(Kwarg)
  12.     def __plot(self, Kwarg):
  13.         save_img_path = os.path.join(self.save_path, 'learning_curve.png')
  14.         plt.clf()
  15.         plt.plot(Kwarg['train_losses'], label='Train', color='g')
  16.         plt.plot(Kwarg['val_losses'], label='Val', color='b')
  17.         plt.xlabel('epoch')
  18.         plt.ylabel('loss')
  19.         plt.legend()
  20.         plt.title('learning_curve')
  21.         plt.savefig(save_img_path)

2.7 main.py

main.py文件将上面所有的东西结合到一起，代码如下：

  1. import torch
  2. import argparse
  4. from model import A2NN
  5. from dataset import Traffic_Light
  6. from utils import get_train_val_names, check_folder
  7. from trainer import Trainer
  8. from validator import Validator
  9. from logger import Logger
  10. from torch.utils.data import DataLoader
  13. def main():
  14.     parse = argparse.ArgumentParser()
  15.     parse.add_argument('--dataset_path', type=str, default='TL_Dataset/')
  16.     parse.add_argument('--remove_names', type=list, default=['README.txt',
  17.                                                              'README.png',
  18.                                                              'Testset'])
  19.     parse.add_argument('--img_resize_shape', type=tuple, default=(32, 32))
  20.     parse.add_argument('--batch_size', type=int, default=1024)
  21.     parse.add_argument('--lr', type=float, default=0.001)
  22.     parse.add_argument('--num_workers', type=int, default=4)
  23.     parse.add_argument('--epochs', type=int, default=200)
  24.     parse.add_argument('--val_size', type=float, default=0.3)
  25.     parse.add_argument('--save_model', type=bool, default=True)
  26.     parse.add_argument('--save_path', type=str, default='logs/')
  28.     args = vars(parse.parse_args())
  30.     check_folder(args['save_path'])
  32.     # pylint: disable=E1101
  33.     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  34.     # pylint: disable=E1101
  36.     model = A2NN().to(device)
  38.     names = get_train_val_names(args['dataset_path'], args['remove_names'])
  40.     train_dataset = Traffic_Light(names['train'], args['img_resize_shape'])
  41.     val_dataset = Traffic_Light(names['val'], args['img_resize_shape'])
  43.     train_dataload = DataLoader(train_dataset,
  44.                                 batch_size=args['batch_size'],
  45.                                 shuffle=True,
  46.                                 num_workers=args['num_workers'])
  48.     val_dataload = DataLoader(val_dataset,
  49.                               batch_size=args['batch_size'],
  50.                               shuffle=True,
  51.                               num_workers=args['num_workers'])
  53.     loss_logger = Logger(args['save_path'])
  55.     logger_dict = {'train_losses': [],
  56.                    'val_losses': []}
  58.     for epoch in range(args['epochs']):
  59.         print('<Main> epoch{}'.format(epoch))
  60.         trainer = Trainer(model, train_dataload, epoch, args['lr'], device)
  61.         train_loss = trainer.train()
  62.         if args['save_model']:
  63.             state = model.state_dict()
  64.             torch.save(state, 'logs/nn_state.t7')
  65.         validator = Validator(model, val_dataload, epoch,
  66.                               device, args['batch_size'])
  67.         val_loss = validator.eval()
  68.         logger_dict['train_losses'].append(train_loss)
  69.         logger_dict['val_losses'].append(val_loss['val_loss'])
  71.         loss_logger.update(logger_dict)
  74. if __name__ == '__main__':
  75.     main()

2.8 compute_prec.py和submit.py

其实上面的七个文件，已经是结束了，下面两个文件一个是用来计算精确度的，一个是用来提交答案的。有兴趣可以看看。
compute_prec.py代码如下：

  1. import torch
  2. import numpy as np
  3. import argparse
  5. from model import A2NN
  6. from dataset import Traffic_Light
  7. from torch.utils.data import DataLoader
  8. from utils import get_train_val_names, check_folder
  11. def main():
  12.     parse = argparse.ArgumentParser()
  13.     parse.add_argument('--dataset_path', type=str, default='TL_Dataset/')
  14.     parse.add_argument('--remove_names', type=list, default=['README.txt',
  15.                                                              'README.png',
  16.                                                              'Testset'])
  17.     parse.add_argument('--img_resize_shape', type=tuple, default=(32, 32))
  18.     parse.add_argument('--num_workers', type=int, default=4)
  19.     parse.add_argument('--val_size', type=float, default=0.3)
  20.     parse.add_argument('--save_path', type=str, default='logs/')
  22.     args = vars(parse.parse_args())
  24.     check_folder(args['save_path'])
  26.     # pylint: disable=E1101
  27.     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  28.     # pylint: disable=E1101
  30.     model = A2NN().to(device)
  31.     model.load_state_dict(torch.load(args['save_path']+'nn_state.t7'))
  33.     model.eval()
  35.     names = get_train_val_names(args['dataset_path'], args['remove_names'])
  37.     val_dataset = Traffic_Light(names['val'], args['img_resize_shape'])
  39.     val_dataload = DataLoader(val_dataset,
  40.                               batch_size=1,
  41.                               num_workers=args['num_workers'])
  43.     count = 0
  44.     for ind, (inp, label) in enumerate(val_dataload):
  45.         inp = inp.float().to(device)
  46.         label = label.long().to(device)
  47.         output = model.forward(inp)
  48.         output = np.argmax(output.to('cpu').detach().numpy(), axis=1)
  49.         label = label.to('cpu').numpy()
  50.         count += 1 if output == label else 0
  52.     print('precision: {}'.format(count/(ind+1)))
  55. if __name__ == "__main__":
  56.     main()

submit.py代码如下：

  1. import torch
  2. import numpy as np
  3. import argparse
  4. import os
  5. import cv2
  7. from model import A2NN
  8. from utils import check_folder
  11. def main():
  12.     parse = argparse.ArgumentParser()
  13.     parse.add_argument('--dataset_path', type=str,
  14.                        default='TL_Dataset/Testset/')
  15.     parse.add_argument('--img_resize_shape', type=tuple, default=(32, 32))
  16.     parse.add_argument('--num_workers', type=int, default=4)
  17.     parse.add_argument('--save_path', type=str, default='logs/')
  19.     args = vars(parse.parse_args())
  21.     check_folder(args['save_path'])
  23.     # pylint: disable=E1101
  24.     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  25.     # pylint: disable=E1101
  27.     model = A2NN().to(device)
  28.     model.load_state_dict(torch.load(args['save_path']+'nn_state.t7'))
  30.     model.eval()
  32.     txt_path = os.path.join(args['save_path'], 'result.txt')
  33.     with open(txt_path, 'w') as f:
  34.         for i in range(20000):
  35.             name = os.path.join(args['dataset_path'], '{}.png'.format(i))
  36.             img = cv2.imread(name)
  37.             img = cv2.resize(img, args['img_resize_shape'])
  38.             img = img.transpose(2, 0, 1)-127.5/127.5
  39.             img = torch.unsqueeze(torch.from_numpy(img).float(), dim=0)
  40.             img = img.to(device)
  41.             output = model.forward(img).to('cpu').detach().numpy()
  42.             img_class = np.argmax(output, axis=1)
  43.             f.write(name.split('/')[2] + ' ' + str(img_class[0]))
  44.             f.write('\n')
  47. if __name__ == "__main__":
  48.     main()

3. 代码如下运行

将数据集下载在文件夹TL_Dataset_Classification，解压后，在TL_Dataset_Classification文件中进入终端，运行命令：

  1. $ python main.py

如果还想计算精确度，在训练玩数据集之后，运行命令：

  1. $ python compute_prec.py

有运行可以到github上提issue或者在给我的邮箱867540289@qq.com发邮件。

4. 结果

学习曲线：
learning_curve.png-22.8kB
在测试集中，实现97.425%的精确度。(继续提升中）

5. 总结

好了，11月的尾巴到此结束，希望能对你学习深度学习问题和pytorch有所帮助。12月马上到，祝我数学考试顺利，也祝各位开开心心！

添加新批注
在作者公开此批注前，只有你和作者可见。
回复批注