В мире искусственного интеллекта мы привыкли к тому, что большие языковые модели (LLM) генерируют тексты за секунды, а компьютерное зрение анализирует фотографии со смартфонов. Однако за пределами цифрового мира существует другая, более масштабная реальность — наблюдение за Землей. Здесь модели работают не с мегабайтами текста, а с терабайтами спутниковых снимков, охватывающими целые континенты. Allen Institute for AI (Ai2) представил платформу OlmoEarth, которая решает одну из самых сложных задач в области геопространственного ИИ: как эффективно запускать инференс (вывод) моделей на планетарном масштабе, не разоряясь и не теряя точности.
OlmoEarth — это семейство фундаментальных моделей, предварительно обученных на примерно 10 терабайтах мультимодальных спутниковых данных. Эти модели уже используются правительствами, НПО и организациями, ориентированными на миссию, для мониторинга вырубки лесов, обеспечения продовольственной безопасности и оценки рисков лесных пожаров. Но сама по себе модель — это лишь половина дела. Для большинства организаций, особенно в экологическом секторе, критическим барьером является не наличие модели, а отсутствие инфраструктуры для её запуска, тонкой настройки и масштабной обработки данных. Именно этот пробел и заполняет OlmoEarth Platform.
01Почему геопространственный инференс — это отдельный вид инженерии?
Большинство ML-моделей принимают на вход несколько мегабайт данных и выдают результат менее чем за секунду. Представьте себе LLM, обрабатывающую абзац текста, или модель компьютерного зрения, анализирующую фото с телефона. Инференс для наблюдения за Землей operates на совершенно другом уровне масштаба. Одна задача по тонкой настройке фундаментальной модели для достижения максимальной производительности может перемещать терабайты данных и выполняться в течение многих часов.
Входы для таких моделей могут охватывать несколько спектральных полос, типы датчиков и временные шаги в пределах большой географической области. Данные могут поступать от нескольких провайдеров, каждый из которых использует разные проекции и разрешения. Кроме того, данные часто содержат пропуски или искажения из-за облачности. Результатом же является не просто текст или класс, а карта, где каждое предсказание должно быть точно выровнено с той же проекцией и координатной сеткой, что и окружающие области.
Даже приобретение данных может стать главной проблемой. Задачи прогнозирования часто тратят больше времени на загрузку и подготовку изображений, чем на сам запуск модели. Это делает эффективные конвейеры данных критически важными. Эти конвейеры должны обрабатывать высокопроизводительный ввод-вывод (I/O), одновременно предоставляя вычислительные мощности для перепроецирования и передискретизации изображений.
02Правильное оборудование для правильной задачи
Поскольку приобретение и подготовка данных часто доминируют во времени выполнения задачи инференса, назначение этой работы на GPU было бы неэффективным. GPU — это самое дорогое оборудование в системе, и оно лучше подходит для параллельных математических вычислений, а не для задач ввода-вывода. Поэтому команда OlmoEarth разделила каждую задачу на три этапа, каждый из которых соответствует определенному профилю оборудования:
- Приобретение и предварительная обработка данных (CPU, высокий I/O): Загрузка, перепроецирование, выравнивание и нормализация изображений, а затем запись их в формате, оптимизированном для быстрой загрузки во время инференса.
- Инференс (GPU): Запуск прямого прохода модели и запись минимально обработанных результатов непосредственно в хранилище.
- Постобработка (CPU): Сшивание результатов по окнам, применение масок или масштабирование, а также экспорт в удобные для пользователя форматы, такие как Zarr, GeoTIFF или GeoJSON.
Платформа OlmoEarth распределяет эти этапы между множеством машин, сохраняя при этом полную загрузку GPU. Многопроцессорные загрузчики данных непрерывно питают каждый GPU, в то время как готовые результаты потоково передаются непосредственно в объектное хранилище (blob storage). Такой подход позволяет избежать простоя дорогостоящих GPU в ожидании данных.
03Один запрос, сотни рабочих процессов и тысячи процессов
Ядро платформы, OlmoEarth Run, является слоем выполнения для задач крупномасштабного инференса. Оно делит географическую область, охваченную каждой задачей, на партиции, размер которых соответствует отдельным вычислительным экземплярам (рабочим узлам). Затем эти партиции subdivируются на меньшие окна, которые модели OlmoEarth обрабатывают по отдельности. Поскольку каждое окно может быть обработано независимо в отдельном прямом проходе, работа над одной частью карты не должна ждать завершения другой.
На практике площадь размером с штат США может стать сотней партиций, а запуск в масштабах континента — тысячами. Соседние партиции немного перекрываются, и мы согласовываем это перекрытие при сборке результатов, чтобы на финальной растре не появлялось швов.
Поскольку партиции независимы, один и тот же этап может выполняться на тысячах вычислительных экземпляров одновременно. Недавно мы использовали этот подход для создания карты риска лесных пожаров, охватывающей всю Северную Америку. В пиковый момент запуск использовал примерно 19 600 процессоров (CPU) и 994 графических процессора (GPU) параллельно, при этом пропускная способность сети превышала 168 ГБ/с. Такой уровень параллелизма сократил расчетное время последовательных вычислений в 4 737 часов до примерно 30,5 часов реального времени — ускорение в 155 раз.
04Поиск и получение нужных пикселей
Заданная географическая область и временной диапазон — это только начало. Платформа должна сначала определить, какие спутниковые сцены должны питать модель. Это означает идентификацию того, что было захвачено, где и когда, у провайдеров с разными каталогами, форматами и задержками публикации. Критерии выбора также зависят от источника: для оптических изображений, таких как Sentinel-2, мы обычно хотим наименее облачные доступные сцены, тогда как для радаров с синтезированной апертурой (SAR) могут быть важнее доступные каналы поляризации.
Мы полагаемся на общедоступные каталоги STAC (SpatioTemporal Asset Catalog) и открытые стандарты, где это возможно. Но крупная задача инференса может генерировать тысячи запросов к метаданным одновременно — гораздо больше, чем внешние сервисы, такие как STAC API Европейского космического агентства (ESA) или Microsoft Planetary Computer, рассчитаны на обработку одновременно.
Чтобы не перегружать эти сервисы, платформа OlmoEarth поддерживает собственный индекс метаданных, который обновляется по мере публикации новых изображений. Для наборов данных, размещенных через AWS Open Data, мы получаем уведомление SNS для каждой новой сцены. Когда провайдер не предоставляет поток изменений, мы опрашиваем его верхний индекс каждые несколько минут. В результате наши запросы к внешним сервисам следуют за стабильным темпом новых публикаций, а не за внезапным всплеском, генерируемым крупной задачей инференса.
Каждая запись индекса хранит метаданные сцены вместе с указателями на каждое место, где доступны исходные пиксели. Во время выполнения платформа выбирает лучший источник и выполняет оконное чтение против форматов, оптимизированных для облака, таких как COG (Cloud Optimized GeoTIFF) или Zarr, извлекая только те байты, которые нужны для данной партиции, а не загружая целые сцены.
Этот индекс также поддерживает наши инструменты аннотации. Поскольку он хранит указатели на изображения Sentinel-1, Sentinel-2, Landsat и NISAR в форматах, оптимизированных для облака, мы можем обслуживать тайлы из любой индексированной сцены через ту же систему оконного чтения, не строя отдельный конвейер ingestion.
05Обработка сбоев на масштабе
Платформа OlmoEarth спроектирована для автоматического восстановления после сбоев. Для каждой задачи в рамках этапа и географической партиции она динамически предоставляет виртуальную машину, работающую с нашим контейнером Docker runner. Runner извлекает параметры задачи, выполняет работу, возвращает результат и завершает работу. Поскольку каждая задача является реентерабельной и идемпотентной, прерывистые сбои можно безопасно обрабатывать путем повторного запуска.
В таком масштабе сбои ожидаемы: провайдер может быть медленным или временно недоступным; метаданные могут указывать на существование изображений, даже когда требуемая полоса или окно отсутствуют; облачность может оставить слишком мало пригодных наблюдений; или задача может полностью аварийно завершиться. Платформа реагирует отслеживанием задач, автоматическими повторными попытками, переходом к альтернативным провайдерам при их наличии и четким различием между ошибками, которые можно повторить, и фатальными ошибками. Отдельный процесс мониторинга обнаруживает зависшие или остановленные runner и перезапускает их задачи.
06Что это значит на практике
Для инженеров и исследователей, работающих с геопространственными данными, опыт OlmoEarth предлагает несколько ключевых уроков. Во-первых, разделение вычислений по типам оборудования (CPU для I/O, GPU для математики) является не просто оптимизацией, а необходимостью для экономии средств и времени. Во-вторых, кэширование и индексация метаданных позволяют избежать перегрузки внешних API и ускорить доступ к данным в разы. В-третьих, идемпотентность задач и автоматическое восстановление от сбоев делают систему устойчивой к нестабильности облачных сред и внешних источников данных.
Платформа OlmoEarth Run требует только виртуальных машин, способных запускать Docker-образ, и доступа к объектному хранилищу. В настоящее время она работает на Google Cloud, но архитектура спроектирована так, чтобы поддерживать несколько облаков и развертывания в собственном аккаунте партнера. Это открывает путь к более демократичному использованию ИИ в экологии и науках о Земле, позволяя организациям с ограниченными ресурсами запускать сложные модели, которые раньше были доступны только крупным корпорациям.
Впереди у команды OlmoEarth есть планы по автоматизации запусков моделей, созданию систем обнаружения изменений с оповещениями, внедрению агентов для упрощения работы с данными, а также развитию моделей встраивания (embeddings) для еще более быстрой и дешевой обработки. Это лишь начало пути, но уже сейчас OlmoEarth демонстрирует, как можно масштабировать ИИ для решения глобальных проблем.
Источник: Hugging Face ↗