В предыдущей статье нам удалось конвертировать custom датасет, доставшийся от yolov5 из формата Pascal VOC в coco формат. Также из датасета Microsoft COCO 2017 (80 классов) были удалены «ненужные» классы, а также присоединены классы из custom датасета (двери и лестницы). В «сухом остатке» 46 классов, из которых 2 вновь добавленных. Осталось дело за малым: обучить модель yolox на вновь созданном датасете, а так же проверить как она работает.
Но для начала yolox необходимо установить.
Установка yolox на windows.
YOLOX использует фреймворк torch, поэтому его необходимо установить заранее, потому как не ясно какую из версий установит YOLOX.
Torch необходим с поддержкой cuda, если видеокарта не от NVIDIA и не поддерживает эту технологию, то обучение модели может растянуться на годы )
Чтобы проверить какая версия драйверов cuda установлена:
nvcc --versionДолжно вывести например:
Версия 10.2. Значит torch устанавливаем:
pip install torch==1.10.0+cu102 torchvision==0.11.0+cu102 torchaudio==0.10.0 -f https://download.pytorch.org/whl/torch_stable.htmlУстановка yolox.
Здесь нет ничего сложного. Скачиваем репозиторий и устанавливаем:
git clone git@github.com:Megvii-BaseDetection/YOLOX.git
cd YOLOX
pip3 install -v -e .
На windows проблем возникнуть не должно, чего не скажешь о других системах, например arm. Но об этом поговорим в другой раз.
Не лишним будет сказать, что у YOLOX есть еще один основной репозиторий с поддержкой MegEngine. Но устанавливается он в разы дольше и набор моделей там не так велик.
Если установка прошла успешно, то можно скачать веса какой-нибудь из моделей yolox и проверить, что модель работает. Также эти веса в дальнейшем понадобятся для обучения.
Проверка yolox.
C сайта репозитория yolox скачаем веса. Лучше выбрать одну из моделей полегче (light), возможно она и не так точна, как ее старшие сестры, однако скорость ее работы, например, на arm в разы больше.
Возьмем картинку и проверим модель:
python tools\demo.py image -n yolox-tiny -c models\yolox_tiny.pth --path assets\bus9.jpg --conf 0.25 --nms 0.45 --tsize 416 --save_resultРезультат найдем в
YOLOX-main\YOLOX_outputs\yolox_tiny\vis_resВремя работы:
Время работы модели (inference time) можно немного уменьшить добавив в строку аргументы --fuse --fp16, а также изменяя параметр --tsize, уменьшив его.
Обучение модели.
Для обучения модели нам понадобятся файлы аннотаций из предыдущей статьи и сам датасет: train и val папки с картинками. Файлы аннотаций — это два файла json, условно train.json и val.json, которые описывают датасет.
Структура директорий перед обучением будет следующая:
/YOLOX/datasets/coco/:
---/annotations/:
------instances_train2017.json
------instances_val2017.json
---train2017 (картинки)
---val2017 (картинки)
Далее необходимо поправить количество классов на которых будет обучаться модель в файле:
YOLOX\yolox\data\datasets\coco_classes.py:
новые классы
COCO_CLASSES = (
"person",
"bicycle",
"car",
"bench",
"cat",
"dog",
"backpack",
"umbrella",
"handbag",
"suitcase",
"sports ball",
"bottle",
"wine glass",
"cup",
"fork",
"knife",
"spoon",
"bowl",
"banana",
"apple",
"sandwich",
"orange",
"carrot",
"chair",
"potted plant",
"bed",
"dining table",
"toilet",
"laptop",
"mouse",
"remote",
"keyboard",
"cell phone",
"microwave",
"oven",
"sink",
"refrigerator",
"book",
"clock",
"vase",
"scissors",
"teddy bear",
"hair drier",
"toothbrush",
"door",
"Stairs"
)
А также необходимо создать файл в папке YOLOX-main\exps\example\custom:
yolo_tiny_my.py
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.
import os
from yolox.exp import Exp as MyExp
class Exp(MyExp):
def __init__(self):
super(Exp, self).__init__()
self.depth = 0.33
self.width = 0.375
self.input_size = (416, 416)
#self.mosaic_scale = (0.5, 1.5)
self.random_size = (10, 20)
self.test_size = (416, 416)
self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
#self.enable_mixup = False
# Define yourself dataset path
self.data_dir = r"D:\3\YOLOX-main\datasets\coco"
self.train_ann = r"D:\3\YOLOX-main\datasets\coco\annotations\instances_train2017.json"
self.val_ann = r"D:\3\YOLOX-main\datasets\coco\annotations\instances_val2017.json"
self.num_classes = 46 #количество классов
self.max_epoch = 300
self.data_num_workers = 4 #4
self.eval_interval = 1
#self.min_lr_ratio = 0.05
Здесь мы указываем количество классов — 46, количество эпох обучения — 300, а также пути, где находится датасет и аннотации.
Запустим процесс обучения.
Будем обучать на NVIDIA GTX 1060 6Gb.
Команда из рабочей директории YOLOX следующая:
python tools/train.py -f exps/example/custom/yolox_tiny_my.py -d 1 -b 8 --fp16 -o -c models/yolox_tiny.pthЕсли вывод команды примерно следующий:
То процесс обучения пошел.
Периодически в консоль будут попадать сообщения:
Параметр ETA показывает приблизительное время окончания обучения.
Осталось подождать 32 дня )
С помощью утилиты windows gpu-z можно посмотреть как трудится видеокарта:
Также с каждой эпохой обучения в YOLOX_Outputs будут попадать сохраненные эпохи:
Это по своей сути модели, которые можно использовать:
python tools/demo.py video -f exps/example/custom/yolox_tiny_my.py -c models/epoch_91_ckpt.pth --path assets/doors1 --conf 0.25 --nms 0.45 --tsize 640 --save_result --device cpuВ примере выше мы прогоняем модель с эпохой обучения 91 через видеоролик и используем только cpu для обработки.
Если вы являетесь счастливым обладателем видеокарты помощнее, например NVIDIA Tesla V100-PCIE-32GB, то процесс обучение модели пойдет значительно бодрее и в строке запуска процесса обучения можно увеличить параметр batch-size (b).
Например так:
python tools/train.py -f exps/example/custom/yolox_tiny_my.py -d 1 -b 128 --fp16 -o -c models/yolox_tiny.pthБлагодарю за внимание.
В планах написать о yolox на raspberry pi и возможностях ускорения модели на arm.
Alex500IS
Спасибо за статью, полезно почитать решения пользователей в этой сфере. Сейчас работаем над обработкой изображений со приборными панелями транспорта и системами учета времени на транспорте
zoldaten Автор
Расскажите, с чем пришлось столкнуться.