Нейросетевой прогноз погоды для edge-устройств

Всем привет! Хочу рассказать не столько про свою модель, сколько про инженерные компромиссы, с которыми я столкнулся во время работы над проектом. Буду рад любой критике.

Речь пойдет о моем проекте edge-weather-forecast — лёгкой нейросетевой модели прогнозирования температуры, которую можно запускать прямо на метеостанции или на простом CPU-устройстве вроде Raspberry Pi.

Почему я вообще решил делать модель для edge-прогноза температуры

Большинство современных ML-моделей для погоды предполагают доступ к глобальным погодным моделям, использование облачных технологий для вычисления на GPU и естественно доступ к Интернету. Для крупных глобальных сервисов это — ок, но не всегда применимо в реальных инженерных сценариях. В некоторых случаях требуется:

• локальный прогноз без доступа к Интернету

• минимальную задержку работы модели

• запуск на слабых CPU-устройствах

Ну например, где это может понадобиться: персональные метеостанции, сельхоз системы мониторинга, удаленные исследовательские точки и вообще любые edge-first сценарии. В этом проекте я сознательно задал жесткие ограничения, чтобы посмотреть, что можно в итоге получить.

Постановка задачи

Итак, вход это история измерений одной станции за последние 672 часа (28 дней): UTC-время, температура, влажность и давление, выход — прогноз абсолютной температуры воздуха на 168 часов (7 дней). Никаких реанализов, гридов и NWP. Также она должна спокойно переноситься на неизвестные при обучении станции.

Далее поговорим о тех самых компромиссах/вопросах и расскажу вообще про свой проект.

Грабля 1. Авторегрессия неудобна для edge

Авторегрессивный прогноз выглядит естественно, на на практике он создает проблемы для edge-сценариев: накопление ошибок, нестабильное поведение и более высокие вычислительные заморочки.

Прямой multi-horizon оказался проще и стабильнее. Один проход модели и все готово. Да, модель стало сложнее обучать, но в целом система становится надежнее.

Почему TCN, а не RNN или Transformer

Пробовал я RNN модели, трансформеры и attention-модели, но довольно быстро от них отказался. Рекуррентные модели очень плохо ведут себя на длинном горизонте, слишком быстро теряют контекст. А трансформеры являются хоть и очень классными для подобной задачи, но слишком тяжелые и по памяти и по времени вычисления, тем более если добавить attention-механизм.

Основой стала всем известная Temporal Convolutional Network (TCN), но с некоторыми доработками. Ее преимущества:

• casual 1-D свертки

• residual

• dilation для большого receptive field

• gated (примерно как в GLU) нелинейности

• фиксированный граф вычислений

• хорошая производительность на CPU

Да, это не самая "модная" архитектура, но в целом очень практичная для моей задачи.

Архитектура модели

Итоговая архитектура довольно простая.

Есть генератор эмбеддинга станций (GNN), он обучается с помощью teacher модели климатологии станций и различать их по расстоянию в км и высоте над уровнем моря (подробнее об это всем можете посмотреть у меня на гитхабе, там есть полный код исследования). Teacher модель отличается от Student тем, что мы с самого старта обучения сливаем часть статических признаков станции в декодер, постепенно все больше зануляя их с каждой эпохой. Это позволяет дать модели горячий старт для более быстрого обучения GNN (различные методы и колбэки можете посмотреть в полном коде моего исследования).

Энкодер и декодер одинаковые для обеих моделей (Teacher и Student), они состоят из нескольких residual gated TCN слоев. Кстати, у меня на гитхабе лежат несколько версий TCN блоков, можете глянуть разные варианты. И в конце легкая выходная голова для генерации финального результата.

Student модель идет уже без GNN внутри себя и имеет меньшее количество слоев и каналов. GNN я вынес в отдельный "продукт", чтобы можно было разово сгенерировать эмбеддинг станции и просто использовать его как константу на инференсе.

Грабля 2. Слишком мощный декодер сломает хвост

Думал я сделать декодер более выразительным, ведь, чем мощнее декодер, тем более "красочный" будет прогноз. Однако, убедился, что это приводит к осцилляциям, локальным провалам и странным скачкам на 3-7 день.

Сделал вывод, что декодер должен быть слабее энкодера. Консервативность здесь важнее гибкости.

Грабля 3. Абсолютная температура усложняет обучение

Предсказывать аномалии в целом проще. Меньше разброс значений и снимает с модели ответственность за сезонность.

Однако, мне пришлось выбрать абсолютную температуру. У нас нет климатологии новой точки, только 28 дней данных. В защиту этого решения приведу еще пару плюсов: выход сразу интерпретируем и инференс из-за этого максимально простой. Цена — более сложное обучение и необходимость следить за формой прогноза.

Грабля 4. MAE плохо отражает качество формы прогноза

Действительно, MAE работает отлично на коротком горизонте. Однако на длинно начинает обманывать. Модель минимизирует MAE, тем самым сглаживает фронты и боится амплитуды. Поэтому в качестве loss метрики я кроме ранжированного MAE использовал большие ошибки, дифференцированную ошибку, экстремальное взвешивание (болье штраф, если модель не предсказала экстремумы) и stddev penalty. Реализацию можно посмотреть на моем гитхабе.

Грабля 5. Амплитуда и метрики конфликтуют

Использование нелинейных блоков (типа GLU) увеличивает амплитуду, делает прогноз "смелее". Однако, MAE становится чуть хуже. Здесь важно было выбрать компромисс: форма прогноза важнее минимальной ошибки.

Грабля 6. EMA весов - не косметика

Изначально при изучении, что такое EMA, оно выглядело как необязательный трюк. Вроде бы, зачем сглаживать веса модели, если и так все хорошо. Однако батчи достаточно нестабильны и веса скачут. EMA сглаживает шум обучения, стабилизирует валидацию и делает прогнозы визуально "человечнее". Действительно, метрика валидации и теста стала намного лучше и стабильнее. EMA-веса также используются в прод-версии модели.

Что в итоге получилось

В результате у меня получилась модель, которая:

• запускается полностью оффлайн

• работает на CPU

• весит ~ 2МБ

• хорошо обобщается на unseen станции

• дает стабильную форму прогноза

Вот такие метрики получились:

взято с моего гитхаба

При этом AP (amplitude preservation) ~ 0.79.

Это не замена глобальным моделям, а локальный инструмент.

Код и материалы

Проект открыт и доступен на GitHub: 👉 edge-weather-forecast

Этот проект стал для меня напоминанием, что инженерные ограничения формируют архитектуру сильнее, чем бенчмарки. Если начинать не с вопроса “как получить лучший скор”, а с вопроса “где и как это будет работать”, многие решения становятся очевидными сами собой.

Буду рад критике, вопросам и идеям.

Источник