使用百度paddlex快速训练图片分类器

---

前言

对内容的分类，有时候不仅需要使用到文本数据，也需要使用到图像数据，将两者联合起来对内容的进行分类。这时候，要是有个工具，能让你快速训练图像分类模型，简直不要太爽，paddlex就是这样一款神器神器。

一、环境依赖

paddlepaddle-gpu或paddlepaddle(版本大于或等于1.8.1)
pip install cython
pip install pycocotools
pip install paddlex==1.3.11

paddlex官方文档

二、使用步骤

1.准备数据

只需要将不同类别的图片放在对应的文件夹内就好了，像这样：

我这里两个类别，分别是live和technology。

里面就是对应的图片

注意： 不能放.gif的图片，会报错

2.划分数据集

命令行进到train_data所在的同级路径中执行：

paddlex --split_dataset --format ImageNet --dataset_dir train_data --val_value 0.2 --test_value 0.1 

执行结果：

3.训练

就几十行代码，真的方便，只需要简单设置训练数据路径，模型保存路径就可以了，请看：

# 设置使用0号GPU卡（如无GPU，执行此代码后会使用CPU训练模型）
import os
import paddlex as pdx
from paddlex.cls import transforms
os.environ['CUDA_VISIBLE_DEVICES'] = '0'


def train():

    train_transforms = transforms.Compose([
        transforms.RandomCrop(crop_size=224),
        transforms.RandomHorizontalFlip(),
        transforms.Normalize()
    ])
    eval_transforms = transforms.Compose([
        transforms.ResizeByShort(short_size=256),
        transforms.CenterCrop(crop_size=224),
        transforms.Normalize()
    ])

    train_dataset = pdx.datasets.ImageNet(
        data_dir='train_data',
        file_list='train_data/train_list.txt',
        label_list='train_data/labels.txt',
        transforms=train_transforms,
        shuffle=True)
    eval_dataset = pdx.datasets.ImageNet(data_dir='train_data',
                                         file_list='train_data/val_list.txt',
                                         label_list='train_data/labels.txt',
                                         shuffle=True,
                                         transforms=eval_transforms)

    num_classes = len(train_dataset.labels)
    model = pdx.cls.MobileNetV3_small_ssld(num_classes=num_classes)

    model.train(num_epochs=20,
                train_dataset=train_dataset,
                train_batch_size=32,
                eval_dataset=eval_dataset,
                lr_decay_epochs=[4, 6, 8, 12, 16, 20],
                save_dir='output/mobilenetv3_small_ssld',
                use_vdl=True)

训练完成后生成如下文件：

4.预测

import numpy as np
import paddlex as pdx
from PIL import Image
from paddlex.cls import transforms

def test():

    predict_transforms = transforms.Compose([
        transforms.RandomCrop(crop_size=224),
        transforms.RandomHorizontalFlip(),
        transforms.Normalize()
    ])
    
    img_path = '0a8d940e2239406ea48a5712edd4ec0b.jpg'

    img = Image.open(img_path)
    img = img.convert("RGB")

    img_cv = np.array(img)

    model = pdx.load_model('output/mobilenetv3_small_ssld/best_model')
    result = model.predict(img_cv)
    # result = model.predict(img_path) 这样也可以
    print(result)

也可以直接给model.predict传图片路径，不过这样更慢，可以自己尝试一下。

总结

paddlex确实方便，很容易上手，效果也不错，我这二分类准确率达到90%+

还可以使用其他很多预训练的模型，有条件的同学可以尝试一下其他模型