基于MAX78000的深度学习垃圾分类器

内容介绍

总结报告

1、项目介绍

随着社会经济的快速发展和群众生活水平的不断提高，生活垃圾产生量持续增加。在这些垃圾中：金属、纸类、塑料、玻璃等是可直接回收利用的资源。但是，在众多种类的垃圾中将有用垃圾回收是较为困难的工作，而且在垃圾分类时也需要耗费巨量的人工。在这样的项目背景下，我们基于MAX78000边缘智能开发板设计了一套深度学习垃圾分类器。通过对垃圾图像进行扫描识别，可以正确分类出有用的垃圾，比如金属、玻璃、纸板等。

2、项目设计思路

2.1 在PC端（ubuntu）搭建开发环境，进行数据集的处理，然后编写模型训练文件，完成模型的训练，量化，评估，进一步生成可以在嵌入式设备上执行的C代码；

2.2 在AI芯片MAX78000上进行模型部署。模型对摄像头采集的图像进行分类识别；

2.3 模型验证的结果分别在TFT-LCD和PC端上位机显示。

X9R9EvAKsF1dwAAAABJRU5ErkJggg==

图 1 设计思路

3、素材搜集思路

本方案的数据集来源于kaggle 的Garbage Classification，该数据集包含来自 6个不同类别的生活垃圾的 2527张图像。其中训练集2000张，测试集527张图片。包含cardboard, glass, metal , paper, plastic and trash共6个分类，其中cardboard, glass, metal , paper, plastic是可回收的生活垃圾。

总垃圾图片

图 2 垃圾数据集

4、项目实现过程

4.1 数据预处理

在ai8x-training/datasets/文件目录下创建garbage.py数据预处理代码文件。在garbage.py中实现了数据集的读入、数据裁剪及数据增强等过程。首先是提取标签，因为此数据集在Kaggle官网下载后，所得到的是已经分好的训练集和测试集图像文件夹，并没有标签，所以需要将图像与对应的标签进行绑定，具体代码如下：

class ImageFolder(Dataset):
    def __init__(self, image_dir, image_size, transform=None):
        # image_size将图像裁剪为指定大小
        self.image_size = image_size
        self.image_paths = []
        self.image_labels = []
        self.transform = transform

        # classes通过listdir得到train/test文件夹下的各个文件夹，用于取出图像和作为标签
        # 打印classes:['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
        self.classes = sorted(os.listdir(image_dir))

        # 通过以下得到图像和标签
        for idx, cls_ in enumerate(self.classes):
            # image_paths每一张图像的路径。
            #如：image_paths:['/home/lyl/MAX78000/ai8x-training/data/garbage/train/cardboard/cardboard236.jpg'
            self.image_paths += glob.glob(os.path.join(image_dir, cls_, '*.*'))
            # image_labels标签
            # [0, 0, 0, 0, 0, 0, 0, 0, 0, 0,..........]
            self.image_labels += [idx] * len(glob.glob(os.path.join(image_dir, cls_, '*.*')))
        # 有多少图像
        self.indexes = list(range(len(self.image_paths)))
.......

接下来是数据增强过程，前面说道此数据集训练集部分共只有2000张图片，这样的数据集规模对于深度学习来说无疑是非常下的，基于此训练的神经网络模型非常容易出现过拟合现象、且模型的学习能力不强。因此对数据进行数据增处理是非常有必要的，经过数据增强后提升模型的学习能力以及模型的泛化能力。具体的数据增强方法有：裁剪—Crop、翻转和旋转—Flip and Rotation、图像变换等方法。在此模型中采用了以下的数据增强方法：

def garbage_get_datasets(data, load_train=True, load_test=True):
    # 数据文件夹位置
    (data_dir, args) = data 
    
    # 要定制的图像大小
    image_size = (64, 64)
    
    # 对训练数据进行预处理
    if load_train:
        data_path =data_dir+'/garbage/train'
        
#         random.shuffle(self.indexes)
        
        train_transform = transforms.Compose([
            transforms.RandomResizedCrop(image_size),
            transforms.RandomRotation(60),
            transforms.RandomGrayscale(p=0.3),
#             transforms.ColorJitter(brightness=1, contrast=1, saturation=1, hue=0.5),
            transforms.ColorJitter(saturation=1, hue=0.5),
            transforms.ToTensor(),
#             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
            ai8x.normalize(args=args)
        ])
        
        train_dataset = ImageFolder(data_path, image_size, train_transform)
        ......

其中transforms.RandomResizedCrop(image_size)是随机长宽比裁剪方法，随机大小，随机长宽比裁剪原始图片，最后将图片裁剪到设定好的image_size。transforms.RandomRotation(30)是随机旋转一定角度，transforms.ColorJitter是修改原始图片的亮度、对比度和饱和度。最后数据集输出格式如下：

datasets = [
    {
        'name': 'garbage',
#         'input': (3, 220, 220),
        'input': (3, 64, 64),
        'output': ('cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash'),
        'loader': garbage_get_datasets,
    },
]

4.2 CNN网络结构设计

通过上述的数据预处理后，明确了输入模型中图像格式为（64x64x64），输出应为6个类别。据此设计卷积神经网络。通过学习在GitHub上美信官方的开发指导（https://github.com/MaximIntegratedAI/ai8x-training），得知MAX78000支持PyTorch开发，但是有些函数是经过改写，并且要基于ai8x。比如FusedMaxPoolCo nv2dReLU==MaxPool2d, followed by Conv2d, and ReLU，因此在设计网络结构时要注意这些变化。首先在/ai8x-training/models/目录下创建ai85net_garbage.py模型文件，此次设计的网络结果共有11层，主要用到了卷积和池化层。首先通过三层连续的卷积层疯狂提取特征。后面几层又通过连续的最大池化和卷积进行图像降维和提取特征。此模型的网络结构以及相关参数如下图：

class Ai85netGarbage(nn.Module):
    """
    SimpleNet v1 Model with BatchNorm
    """
    def __init__(
            self,
            num_classes=6,
            num_channels=3,
            dimensions=(64, 64),  # pylint: disable=unused-argument
            bias=True,
            **kwargs
    ):
        super().__init__()
        # (64,64)
        self.conv1_1 = ai8x.FusedConv2dBNReLU(num_channels, 64, 3, stride=1, padding=1, bias=True, batchnorm='NoAffine', **kwargs)
        # (64,64)
        self.conv1_2 = ai8x.FusedConv2dBNReLU(64, 32, 1, stride=1, padding=0, bias=True, batchnorm='NoAffine', **kwargs)
        # (64,64)
        self.conv1_3 = ai8x.FusedConv2dBNReLU(32, 64, 3, stride=1, padding=1, bias=True, batchnorm='NoAffine', **kwargs)
         # (32,32)
        self.conv2_1 = ai8x.FusedMaxPoolConv2dBNReLU(64, 32, 3, stride=1, padding=1, bias=True, batchnorm='NoAffine', **kwargs)
        
        self.conv2_2 = ai8x.FusedConv2dBNReLU(32, 64, 1, stride=1, padding=0, bias=True, batchnorm='NoAffine', **kwargs)
        # (16,16)
        self.conv3_1 = ai8x.FusedMaxPoolConv2dBNReLU(64, 128, 3, stride=1, padding=1, bias=True, batchnorm='NoAffine', **kwargs)
        
        self.conv3_2 = ai8x.FusedConv2dBNReLU(128, 128, 1, stride=1, padding=0, bias=True, batchnorm='NoAffine', **kwargs)
        # (8,8)
        self.conv4_1 = ai8x.FusedMaxPoolConv2dBNReLU(128, 64, 3, stride=1, padding=1, bias=True, batchnorm='NoAffine', **kwargs)
        
        self.conv4_2 = ai8x.FusedConv2dBNReLU(64, 128, 3, stride=1, padding=1, bias=True, batchnorm='NoAffine', **kwargs)
        # (4,4)
        self.conv5_1 = ai8x.FusedMaxPoolConv2dBNReLU(128, 128, 1, stride=1, padding=0, bias=True, batchnorm='NoAffine', **kwargs)
        
        # (2,2)
        self.conv5_2 = ai8x.FusedMaxPoolConv2dBNReLU(128, 128, 1, stride=1, padding=0, bias=True, batchnorm='NoAffine', **kwargs)
        
         self.fc = ai8x.Linear(512, num_classes, bias=True, wide=True, **kwargs)
......

在设计网络结构以及参数时要注意的是，卷积核大小仅仅可以设置为3和1，在复现经典网络时，要注意更改卷积核大小和其他相关参数。

4.3 模型训练、量化以及生成c代码过程

（a）模型训练

在写好数据预处理代码和模型训练代码后，就可以进行模型的训练了，在/ai8x-training/models/scripts/目录下创建train_garbage.sh模型训练脚本，脚本内容为：

#!/bin/sh
python train.py --deterministic --epochs 500 --optimizer Adam --lr 0.00018 --wd 0 --compress policies/schedule-cifar-nas.yaml --model ai85net_garbage --dataset garbage --device MAX78000 --batch-size 16 --print-freq 100 --validation-split 0 --use-bias --qat-policy policies/qat_policy_late_cifar.yaml --confusion "$@"

在其中指定了训练的epoch 500，学习率lr 0.00018，其中优化方法选择Adam，batch-size为16，这些参数都是经过多次训练后，所调整的一些使模型学习能力较优的参数。其次如果在量化时，选择量化方法为QAT（量化感知训练），要设置--qat-policy，并编写qat_policy_late_cifar.yaml文件，还有其他的一些设置参数，比如指定开发板--device MAX78000，在模型训练文件中使用偏差bias要指定--use-bias，否则会报错等这些注意事项，通过阅读官方的开发指南都可得到这些信息。

FtUQn7lQIKj_dQmHzB7wd13DsxLF

（b）模型量化

训练好模型之后，便会得到模型参数文件，在ai8x-training/models/logs文件夹中会生成以训练时间为命名的文件夹，其中包含了模型的训练日志，以及模型文件等。需要将效果最好的模型文件qta_best.pth.tar拷贝到/ai8x-synthesis/trained/，并在/ai8x-synthesis/scripts下创建量化脚本命令：

#!/bin/sh
python quantize.py trained/qat_garbage.pth.tar trained/qat_garbage-q.pth.tar --device MAX78000 -v "$@"

（c）模型评估

得到量化后的模型之后，在ai8x-training/scripts/目录下创建模型评估脚本命令，目的是为了评估模型量化后的识别能力。脚本命令为：

#!/bin/sh
python train.py --batch-size 16 --optimizer Adam --lr 0.00018 --model ai85net_garbage --dataset garbage --confusion --evaluate --exp-load-weights-from ../ai8x-synthesis/trained/qat_garbage-q.pth.tar --device MAX78000 -8 --use-bias "$@"

（d）生成样本测试文件

在ai8x-training/test_samle/目录下创建测试文件，生成的文件拷贝到ai8x-synthesis/test/目录下备用。脚本文件命令为：

./train.py --batch-size 16 --optimizer Adam --lr 0.00018 --model ai85net_garbage --save-sample 10 --dataset garbage --evaluate --exp-load-weights-from ../ai8x-synthesis/trained/qat_garbage-q.pth.tar -8 --device MAX78000 "$@"

（e）生成开发板所能执行的代码

首先要在ai8x-synthesis/networks下创建网络结构yaml文件，然后在ai8x-synthesis/transform_c下创建转化模型脚本命令文件，在脚本命令中可以指定是否使用fifo，以及模型生成的文件夹等信息，并调用编写好的网络结构yaml文件。

yaml文件

arch: ai85net_garbage
dataset: garbage

layers:
  - out_offset: 0x4000
    processors: 0x0000000000000007  # 1_1
    operation: conv2d
    kernel_size: 3x3
    pad: 1
    activate: ReLU
    data_format: HWC
  - out_offset: 0x0000
    processors: 0xffffffffffffffff  # 1_2
    operation: conv2d
    kernel_size: 1x1
    pad: 0
    activate: ReLU
  - out_offset: 0x4000
    processors: 0x00000000ffffffff  # 1_3
    operation: conv2d
    kernel_size: 3x3
    pad: 1
    activate: ReLU
......

脚本命令

./ai8xize.py --verbose --test-dir demos --prefix ai85-garbage_cifar74_no_fifo --checkpoint-file trained/qat_garbage-q.pth.tar --config-file networks/garbage_cifar.yaml --device MAX78000 --compact-data --mexpress --softmax --overwrite

Fqfjgl3h4I2DjSNknszqmnCbnR8d