Описание идеи
Молнии очень непредсказуемое природное явление. Просто так ты никогда не сможешь определить, куда и когда она ударит. Но самое непредсказуемое — это её форма. Изгибы, которые как снежинки отличаются одна от другой.
Принцип работы кода
- Монтирование Google Drive Зачем: Чтобы иметь доступ к своим файлам и папкам, хранящимся в Google Drive, и сохранять результаты обучения там. Что я сделала:
Проверила, подключён ли диск по пути /content/drive/MyDrive. Если не подключён — вызвала функцию drive.mount() для подключения.
- Определение конфигурации обучения Зачем: Чтобы централизованно настроить параметры обучения, такие как пути к данным, размер изображений, гиперпараметры, моделевые параметры и др. Что я сделала:
Создала словарь CONFIG, в котором указала все нужные параметры, например: путь к папке с изображениями, название модели, размеры, скорость обучения, количество эпох и других.
- Определение устройства и типа данных Зачем: Чтобы использовать GPU для быстрого обучения, если он доступен, и выбрать правильный тип данных для оптимизации. Что я сделала:
Проверила наличие cuda (torch.device('cuda')) или использовала CPU (torch.device('cpu')). В зависимости от этого выбрала тип данных (float16 для cuda для ускорения и снижения потребления памяти, float32 для cpu). Установила сид (seed) для воспроизводимости обучения.
- Создание функции gather_images() Зачем: Чтобы собрать все пути к изображениям внутри заданной папки и её подпапок, чтобы позже использовать их для загрузки данных. Что я сделала:
Прошлась по папке с помощью os.walk(), собрала все файлы с допустимыми расширениями (.png, .jpg, .jpeg, .webp), чтобы иметь список всех изображений.
- Создала класс LightningDataset Зачем: Чтобы удобно управлять датасетом изображений и их текстовых подсказок, а также подготовить их для подачи в модель. Что я сделала:
В конструкторе класса: Нашла все пути к изображениям в папке. Перемешала список файлов произвольно для рандомизации. Ограничила общее количество изображений по лимиту (limit). Создала серию преобразований изображений — изменение размера, преобразование в тензор и нормализация. В методе getitem(): Загружала изображение. Применяла преобразования. Возвращала словарь из изображения и текста.
- Подготовка к обучению: создание датасета и DataLoader Зачем: Чтобы обеспечить эффективную подачу данных в модель по батчам с возможностью перетасовки. Что я сделала:
Создала объект LightningDataset, указав папку с изображениями. Создала DataLoader, который подает батчи изображений и подсказок во время тренировки.
- Загрузка предобученной модели Зачем: Чтобы использовать уже обученную модель Stable Diffusion в качестве основы, на которую будут добавляться настройки LoRA. Что я сделала:
Загрузила модель через DiffusionPipeline, указав название нужной предобученной модели. Отключила встроенные проверки безопасности — чтобы ускорить обучение. Переместила модель на выбранное устройство (cuda или cpu).
- Подготовка модели для обучения LoRA Зачем: Чтобы сделать модель способной адаптироваться к новым данным, сохраняя при этом основные параметры. Что я сделала:
Взяла unet из пайплайна — основную часть модели диффузии. Создала конфигурацию LoRA с помощью LoraConfig, указав параметры размерности, имя, масштаб и др. Обернула unet в get_peft_model(), чтобы применить LoRA — это добавит обучаемые низкоранговые матрицы.
- Настройка оптимизатора Зачем: Чтобы обновлять параметры модели во время обучения. Что я сделала:
Создала AdamW для параметров модели (model.parameters()), установив скорость обучения.
- Начало цикла обучения Зачем: Чтобы пройтись по всем обучающим данным, обновляя веса модели при помощи градиентных методов. Что я сделала:
Задала переменные для отслеживания общего количества шагов, времени обучения и прогресса.
- Обработка каждого батча Зачем: Чтобы подготовить входные данные, провести обучение, посчитать ошибку и обновить модель. Что я сделала:
Передала изображения на устройство. Преобразовала их в латентное пространство через VAE (encode(), latent_dist.sample()) и масштабировала их (* 0.18215). Токенизировала и закодировала подсказки (tokenize() и encode()), чтобы получить эмбеддинги текста. Создала шум (noise) и выбрала случайное время t для процесса диффузии (sample_t()). Внесла шум в латенты с помощью add_noise(). Передала зашумленные латенты и временные параметры в модель — она предсказывает шум (model(noisy_latent, timesteps=t)). Посчитала потерю как разницу между предсказанным и реальным шумом (MSE). Обнулила градиенты (optimizer.zero_grad()), произвела обратное распространение (loss.backward()), и обновила параметры (optimizer.step()).
- Сохранение контрольных точек Зачем: Чтобы иметь возможность вернуться к успешной точке обучения. Что я сделала:
Каждые save_interval шагов создавала папку и сохраняла веса модели (save_pretrained()).
- Завершение обучения Зачем: Чтобы сохранить финальную версию модели. Что я сделала:
После достижения max_steps — создала папку final и сохранила финальные веса модели.
