Можно ли обучать upscale-модель на 4 ГБ VRAM: короткий ответ
Да, обучение модели суперразрешения на 4 ГБ видеопамяти возможно, но каждый шаг упирается в ограничения железа. Команда Neoflex обучила upscale-модель для реставрации старых фото и видео на ноутбуке MSI Katana GF66 с видеокартой RTX 3050 Ti Laptop (4 ГБ VRAM). Рабочие параметры: batch_size=6, обучение в смешанной точности (AMP), около 60 часов на одну эпоху. Метрика PSNR выросла с 14.3 до 16.6 дБ за 4 эпохи. Описание проекта и пайплайна авторы опубликовали на Habr.
Цена запуска на слабой GPU: batch_size=6 вместо привычных 32-64, десятки часов на эпоху, ручная очистка памяти и отключение cuDNN benchmark. Без этих мер обучение падает с ошибкой CUDA out of memory.
Результат 16.6 дБ PSNR стоит воспринимать трезво. Это низкий показатель: изображения остаются размытыми, часть мелких деталей теряется. На этом этапе модель доказывает работоспособность процесса, до продакшена ей далеко. Опыт полезен тем, кто хочет разобраться в механике обучения image-to-image моделей и увидеть реальные границы слабого железа.
Зачем Neoflex взялись за обучение своей модели реставрации
Цель проекта: автоматизировать реставрацию старых материалов, сохранив их аутентичность. Готовые сервисы часто добавляют избыточную резкость или некорректно обрабатывают характерные дефекты плёнки, например царапины и цветовой сдвиг.
Второй мотив технический. Чтобы разобраться в механике обучения и инференса, нужен полный доступ к весам и архитектуре. Закрытые API такого доступа не дают, поэтому авторы собрали собственный пайплайн.
Типичные сценарии, под которые писалась модель: оцифровка семейных видеокассет и архива институтских плёнок. Объёмы личных и ведомственных архивов большие, а требования к точности цветопередачи высокие. Проект задуман как исследовательский.
Постановка задачи: что именно обучает модель
Формально задача сформулирована как supervised image-to-image translation, то есть обучение с учителем для преобразования одного изображения в другое. Восстановление ретрофотографий и кадров очень низкого качества сводится к паре «вход-эталон».
- Вход: X_LR ∈ R^{H×W×3}, изображение низкого разрешения с наложенным синтетическим шумом.
- Выход: Y_HR ∈ R^{2H×2W×3}, изображение высокого разрешения, максимально близкое к эталону. Коэффициент увеличения - 2x.
Функция потерь комбинированная: L1 (MAE) отвечает за попиксельную точность, perceptual loss на признаках предобученной VGG - за визуальное качество. Качество измеряют метрикой PSNR (пиковое отношение сигнал/шум). Формулировка задачи и набор лоссов приведены в исходной статье.
Почему L1, а не L2. Квадратичная ошибка сильнее штрафует крупные отклонения и подталкивает модель к усреднению, что даёт размытые изображения. L1 ведёт себя мягче и меньше «замыливает» детали. Perceptual loss работает с признаками на промежуточных слоях VGG: расхождение в текстурах и структуре получает вес, которого у попиксельных метрик нет.
Архитектура на Residual Blocks: почему именно она
Основа модели - Residual Blocks. Идея простая: блок предсказывает поправку, которую нужно добавить к входу, чтобы получить выход. Остаточные связи (skip connections) пропускают сигнал в обход слоёв, поэтому градиент доходит до ранних слоёв без затухания, а сеть можно делать глубже без потери качества обучения.
Для 4 ГБ VRAM это ещё и вопрос памяти: остаточные связи дают активациям короткие пути, и часть промежуточных тензоров переиспользуется вместо хранения до конца прямого прохода. Точную глубину сети и число блоков в доступном фрагменте статьи авторы не раскрывают.
Настройка окружения на Windows: CUDA 13.0, драйверы NVIDIA и uv
Авторы работали на Windows и собрали окружение из трёх частей: свежий драйвер NVIDIA, CUDA Toolkit 13.0 и виртуальное окружение через uv. Порядок такой:
- Обновить драйвер NVIDIA до версии, поддерживающей нужную ветку CUDA.
- Установить CUDA Toolkit 13.0.
- Создать виртуальное окружение через uv.
- Поставить в него сборку PyTorch, собранную под CUDA 13.0.
Главное правило: драйвер, CUDA Toolkit, cuDNN и сборка PyTorch должны совпадать по ветке. Если PyTorch собран под CUDA 12.x, а в системе стоит 13.0, torch.cuda.is_available() вернёт False либо обучение упадёт на первом же тензорном ядре. Проверять совместимость стоит до запуска модели, чтобы не ловить ошибку на первом шаге.
uv выбран вместо pip: он быстрее ставит пакеты и строже разрешает зависимости, что на Windows заметно экономит время. Точные команды установки в доступном описании не приводятся, поэтому сверяйтесь с официальными инструкциями PyTorch и NVIDIA под свою версию.
Датасет SRDataset: четыре типа синтетического шума
Данные собраны в кастомный датасет SRDataset из 12 000 объектов. Пары «низкое разрешение - высокое разрешение» получают синтетически: берут чистый кадр, накладывают шум и получают вход, а исходник служит эталоном. Так обходятся без ручной разметки.
| Тип шума | Что имитирует | Зачем в датасете |
|---|---|---|
| Гауссов шум | Зернистость плёнки, шум цифрового сенсора при слабом свете | Модель учится отличать случайные отклонения яркости от реальных деталей |
| «Соль и перец» | Битые пиксели, выпадения при оцифровке | Убирать одиночные чёрные и белые точки, не размазывая окружение |
| Квантование цвета | Потеря глубины цвета на старых носителях | Восстанавливать плавные переходы там, где осталось несколько оттенков |
| Цветной «соль и перец» | Дефекты цветных плёнок, сбои каналов | Гасить цветные артефакты, не превращая кадр в чёрно-белый |
Пропорции смешивания шумов и силу каждого из них в доступном описании авторы не указали. Если повторять пайплайн, эти параметры подбирают под свой архив: чем сильнее зашумлён вход, тем дольше модель выходит на приемлемый PSNR.
Обучение в смешанной точности: batch_size=6 и 60 часов на эпоху
Обучение шло в смешанной точности (AMP) с batch_size=6 на 12 000 объектов. Одна эпоха занимала около 60 часов, четыре эпохи - примерно 240 часов, то есть около 10 суток непрерывной работы ноутбука.
AMP держит веса в float32, а вычисления на свёртках и матричных умножениях ведёт в float16 там, где это безопасно по диапазону значений. Выигрыш двойной: активации занимают меньше памяти и тензорные ядра работают быстрее. На 4 ГБ VRAM без AMP модель просто не поместилась бы.
batch_size=6 - очень мало даже по меркам одной GPU. Такой размер даёт шумные градиенты, из-за чего сходимость медленнее, а слой нормализации считает статистику всего по шести примерам за шаг. Дальнейший рост batch_size упирается в память: на 4 ГБ запас заканчивается почти сразу.
60 часов на эпоху объясняются слабым железом. Архитектура модели здесь ни при чём. Замеров на других конфигурациях в источнике нет, поэтому сравнивать время обучения с другими GPU в цифрах не будем.
Борьба с ошибками CUDA out of memory и CUDNN_STATUS_EXECUTION_FAILED
Две ошибки преследовали авторов на всём протяжении обучения. CUDA out of memory возникает, когда PyTorch не может выделить очередной тензор: память занимают активации, кэш аллокатора и градиенты. CUDNN_STATUS_EXECUTION_FAILED - отказ библиотеки cuDNN запустить выбранный алгоритм свёртки.
Приёмы, описанные в опыте Neoflex: отключение cuDNN benchmark, включение детерминизма, ручная очистка памяти и мониторинг температуры памяти.
- Отключить cuDNN benchmark.
- Включить детерминированные алгоритмы.
- Чистить память вручную: gc.collect() и torch.cuda.empty_cache().
- Следить за температурой Memory Junction.
Почему cuDNN benchmark и детерминизм важны на 4 ГБ VRAM
Флаг cudnn.benchmark=True заставляет cuDNN при первом проходе перебрать доступные алгоритмы свёртки под конкретный размер входа и выбрать самый быстрый. Перебор требует дополнительной памяти и времени, а на слабой GPU часть алгоритмов может не запуститься, что и даёт CUDNN_STATUS_EXECUTION_FAILED. Отключение benchmark убирает подбор, снижает пиковое потребление памяти и делает поведение предсказуемым. Плата - возможная потеря скорости на постоянном размере входа, где лучший алгоритм всё же нашёлся бы.
Детерминизм (torch.use_deterministic_algorithms(True)) запрещает недетерминированные операции: один и тот же код даёт один и тот же результат при повторном запуске. Побочный эффект: часть быстрых алгоритмов cuDNN исключается из выбора, обучение может идти медленнее. На 4 ГБ VRAM разумно жертвовать скоростью ради воспроизводимости и стабильности.
Ручная очистка памяти: gc.collect() и torch.cuda.empty_cache()
PyTorch не возвращает освободившиеся блоки памяти драйверу сразу: их держит кэш аллокатора, чтобы не запрашивать память заново. При жёстком лимите в 4 ГБ этот кэш мешает. gc.collect() запускает сборщик мусора Python и убирает объекты, на которые больше нет ссылок, например тензоры из предыдущих итераций. torch.cuda.empty_cache() отдаёт неиспользуемые блоки из кэша CUDA.
Вызывать обе функции стоит, когда ссылки на тензоры потеряны, но память не вернулась: после эпохи или после шага валидации. Частые вызовы замедляют обучение, потому что каждый следующий шаг снова запрашивает память у драйвера. Конкретный интервал авторы не указывают, его подбирают по показаниям nvidia-smi.
Мониторинг температуры Memory Junction
Memory Junction - температура микросхем памяти GPU. На ноутбуке с плотной компоновкой она растёт быстрее, чем температура ядра, а троттлинг памяти бьёт по скорости обучения незаметно: ошибок нет, просто эпоха идёт дольше. Смотреть показатель удобно через nvidia-smi и утилиты мониторинга, на Windows подойдут HWiNFO и MSI Afterburner.
Если температура держится у верхней границы, вариантов несколько: снизить нагрузку, улучшить охлаждение (подставка, чистка системы от пыли, андервольт), сократить длину непрерывного прогона. Точные пороги авторы не приводят, значения зависят от модели ноутбука.
Результаты: рост PSNR с 14.3 до 16.6 дБ за 4 эпохи
За четыре эпохи PSNR вырос с 14.3 до 16.6 дБ, то есть на 2.3 дБ. Разница в 2.3 дБ соответствует снижению среднеквадратичной ошибки примерно в 1.7 раза: динамика есть, градиенты не затухли, выбранные лоссы работают.
Абсолютное значение остаётся низким. Картинка становится чище, но до реставрационного качества далеко: на таких значениях PSNR заметны размытость и артефакты. Обучение на этом не закончилось, и модель не готова к продакшену. Четыре эпохи - это разведка, а не финальный результат.
Практическая ценность опыта в другом. Он показывает, как выглядит полный цикл обучения SR-модели на слабой GPU, какие настройки обязательны, с какой скоростью растёт метрика и сколько времени реально занимает итерация.
Кому подходит такой подход и когда лучше выбрать облако
Опыт пригодится, если вы хотите понять механику обучения image-to-image моделей целиком и готовы ждать. 60 часов на эпоху означают, что цикл «поменял гиперпараметр - посмотрел результат» занимает дни. Такой режим подходит исследователям и энтузиастам, у которых есть архив данных и нет жёсткого дедлайна.
Если нужен быстрый результат или высокое качество, облачные GPU обойдутся дешевле времени. Аренда ускоряет итерации в разы, а готовые API закрывают типовые задачи реставрации без обучения. Ограничение облака известно: доступ к весам и архитектуре вы не получаете, а именно этого не хватало авторам.
Инференс - отдельная история. Готовую модель реально запускать на том же ноутбуке, потому что он требует памяти в разы меньше, чем обучение с градиентами и состояниями оптимизатора. Тяжело даётся именно цикл обучения.
Есть и другие способы учить модели на ограниченном железе: LoRA-обучение DeepSeek-V4-Flash с GGUF и кастомными Triton-ядрами укладывает модель в 90 ГБ VRAM, а fine-tuning Stable Diffusion на CPU Intel показывает, как обойтись вообще без GPU. Для реставрации фото это смежные задачи, но приём общий: сокращать память под градиенты и состояния оптимизатора.
Если повторять пайплайн Neoflex, начинайте с малого: несколько сотен изображений, batch_size=2-4, одна эпоха для проверки окружения на ошибки. Растить датасет имеет смысл только после того, как эпоха проходит без CUDA out of memory.