Python version

3.10.14

Финальная модель

В качестве финальной модели я выбрал unet-largest-tw(жёлтый). Вы также можете использовать unet-large(фиолетовый), которая не сильно проседает по метрикам и картинкам, зато обучение и inference у нее быстрее.

Запуск inference

• unet-largest-tw (желтый): python3 inference.py exp.config_path=configs/unet-largest-tw-test-final.yaml train.checkpoint_path=<путь к чекпоинту>, ссылка на чекпоинт
• unet-large (фиолетовый): python3 inference.py exp.config_path=configs/unet-large-test-final.yaml train.checkpoint_path=<путь к чекпоинту>, ссылка на чекпоинт

Результатом исполения являются 2020 сгенерированных картинок в папке experiments/images, а также файл со значением тестовой метрики FID experiments/test_metrics.txt. Значение метрики также выводится в stdout.

Советы

• Предолагается, что датасет лежит в директории проекта, то есть в той же где и inference.py. Если это не так, добавьте к строке запуска или измените в файле конфига параметры:
  • data.dataset_dir - директория датасета, где находятся папки train и test
  • data.input_val_dir - папка test - в ней папки тестовых классов
• Предполагается что в системе доступно 16ГБ видеопамяти, если это не так можно уменьшать data.val_batch_size в конфигах
• Предполагается что директория откуда запущен inference доступна для записи сгенерированных изображений
• Можно залогировать в wandb 25 картинок и тестовую метрику. Для этого нужно задать переменную среды WANDB_KEY, и добавить в конфиг exp.use_wandb=True, wand.project=<ваш проект>, опционально wandb.run_name=<имя рана>

Запуск train

• unet-largest-tw (жёлтый):
  1. python3 train.py exp.config_path=configs/unet-largest-tw-train-final.yaml
  2. Сохранить чекпоинт, он будет использоваться для дообучения.
  3. python3 train.py exp.config_path=configs/unet-largest-tw-finetune-final.yaml
• unet-large (фиолетовый): python3 train.py exp.config_path=configs/unet-large-train-final.yaml

Результат исполнения - чекпоинт experiments/checkpoints/checkpoint-latest.pth

Советы

• Опять же, предполагается что датасет лежит там же где и train.py. Можно изменить в конфиге:
  • data.dataset_dir - директория датасета, где находятся папки train и test
  • data.input_train_dir - папка train - в ней папки с обучающими изображениями для каждого класса
• Предполагается что директория откуда запускается обучение доступна для записи чекпоинта
• Предполагается что в системе доступно 16ГБ видеопамяти, если это не так можно уменьшать data.train_batch_size и data.val_batch_size в конфигах
• По умолчанию обучение логирует только tqdm полоску с общим прогрессом, и использует весь датасет. Можно обучать модель так как это делал я, для этого:
  • Использовать модифицированный датасет, в котором от трейна отделена валидация аналогичная тесту
  • Тогда если датасет не там же где и train.py, еще нужно указать data.input_val_dir=<директория с папками валидационных изображений>
  • Настроить логирование в wandb. В переменной среды WANDB_KEY указать свой ключ, добавить в конфиг exp.use_wandb=True, wand.project=<ваш проект>, опционально wandb.run_name=<имя рана>. В wandb будут все метрики и картинки.
  • Параметр train.val_images_per_class установлен на 6, т.е. на валидации оставляем 6 изображений от каждого класса. Можно увеличить если готовы ждать