За один вечер, в домашних условиях, без регистрации, смс и всего такого 😉
Без облаков, без AWS и без ML-кластера - только твой ноутбук.
Будем разбирать обучение на примере личных машин: ПК на windows c RTX 4070 на борту и MacBook Pro на обычном M4
Что будет в этом посте:
- Поднимем окружение для обучения YOLO на Windows и macOS
- Проверим работу предобученной модели
- Скачаем датасет и запустим своё обучение
- Получим локальную модель обученную под наши цели
Если есть Macbook или приемлемая видеокарта и поставить Python - повторить можно уже сегодня.
Даже без топовой видеокарты - просто будет дольше.
Контекст задачи
Напомню, что в универе в рамках курса для студентов стояла задача реализовать computer vision модель по определению СИЗ на опасных производствах.
На первом этапе, год назад, мы делали прототип - обычное веб приложение с выходом на готовое API и определением СИЗ на отправляемых изображениях.
В этом году решили реализовать следующий этап - обучить локальную модель на готовом датасете.
Один из проактивных студентов, Саша, взял на себя ресерч и обучение на стороне Windows, за что ему огромная благодарность.
Если кому интересно самому обучить модель под свои цели, то предоставлю шаги и инструкцию как это можно сделать:
Технический этап - установить пакеты и подготовить окружение:
Нам понадобится Python, пакетный менеджер pip и библиотека ultralytics (YOLO).
Устанавливаем Python для своей платформы win/mac из python.org/downloads
Рекомендую стабильного старичка версии 3.10
После установки можно запустить файл IDLE и поиграться или проверить в терминале:
python --version
pip - менеджер пакетов, а-ля фронтовый npm
pip3 install --upgrade pip
На маке использовал версию pip3, на винде предыдущую
pip install --upgrade pip
ultralytics - библиотека для задач компьютерного зрения от одноименной компании, нужна для работы с популярной моделью YOLO по обнаружению объектов на изображениях.
pip3 install ultralytics # для винды без 3
Стоит обратить внимание на сообщения в терминале, особенно про PATH.
Если что‑то непонятно - скормить команду и вывод в нейросеть и спросить: «Я выполнил команду X, в терминале отобразилось Y, всё ок или надо что‑то исправить?».
На этом этапе у тебя уже:
- Готов Python
- Обновлён pip
- Установлена библиотека ultralytics
- Можно запускать предобученные модели
Запуск и проверка
В базовой модели около 70 объектов для распознавания. Хватит для предметов домашнего обихода )
Я запускал как легкую модель yolov8n так и большую yolov8l - разницы не заметил, обе прекрасно работают на первый взгляд.
Проверим что всё работает:
Пример кода для запуска базовой модели закину в комменты. В корне проекта создаем любой файл с расширением .py и закидываем код туда. Запуск командой:
python3 file.py
PyTorch - фреймворк, который нам понадобится для обучения своей модели.
Установка для macOS:
pip3 install torch torchvision torchaudio
Для Windows:
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
Для винды указываем конкретную версию пакета с CUDA для работы с видеокартой, обучение на мощностях CPU дольше во много раз.
В маке установка стандартная, а использование графического процессора GPU Apple (MPS) указываем в скрипте запуска обучения.
После установки проверим что всё ок
Скрипт проверки возможности обучения на gpu закину вторым комментом
Если torch установился без ошибок - мы уже технически готовы к обучению своей модели.
Дальше самое интересное - работа с данными.
Качаем датасет
Для задачи с СИЗ мы использовали готовый датасет по PPE (каски, жилеты, перчатки и др.) с Roboflow.
Ищем:
https://universe.roboflow.com/search?q=ppe
Выбираем датасет. Есть как чисто датасет, так и с обученными моделями. Делаем форк в свой воркспейс и скачиваем в формате для YOLO себе на устройство и распаковываем.
На этом этапе у тебя:
- Готовое рабочее окружение
- Понимание как запускать модель
- Скачанный датасет в формате YOLO
Остался только основной этап с обучением своей модели, но он не влез в один пост 😁 ⬇️
1/2