La plataforma OlmoEarth: Inferencia geoespacial a escala planetaria

🌍 Obtenga más información sobre la plataforma OlmoEarth: allenai.org/olmoearth](https://allenai.org/olmoearth)

Los modelos OlmoEarth conforman nuestra familia de modelos fundacionales de observación terrestre, preentrenados con aproximadamente 10 terabytes de datos satelitales multimodales. Gobiernos, ONG y otras organizaciones orientadas a objetivos específicos ya están adaptando OlmoEarth para aplicaciones que incluyen el monitoreo de la deforestación, la seguridad alimentaria y el riesgo de incendios forestales.

En Ai2, sabemos cómo entrenar y liberar modelos abiertos potentes, y para organizaciones con equipos de ingeniería sólidos, un modelo abierto es todo lo que necesitan para comenzar. Pero la mayoría de las organizaciones en el ámbito ambiental —las que están mejor posicionadas para aplicar estos modelos— no poseen la infraestructura ni los equipos de ingeniería capaces de gestionar el ciclo de vida completo: etiquetado de datos, ajuste fino (fine-tuning) y ejecución de inferencia a gran escala. Hemos dedicado más de una década a operar plataformas como Skylight y EarthRanger, software en el que los usuarios de todo el mundo confían a diario, por lo que debe funcionar todos los días. Esa experiencia nos enseñó lo que requiere generar impacto: ejecutar modelos de manera rentable en el momento y lugar adecuados, monitorear el rendimiento, convertir resultados brutos en información accionable y verificar que dichos resultados impulsen los objetivos que nuestros socios desean.

Es por eso que construimos la Plataforma OlmoEarth: infraestructura para llevar modelos geoespaciales desde el ajuste fino y la evaluación hasta la inferencia a gran escala.

La inferencia a esta escala presenta su propio conjunto de desafíos. Las imágenes satelitales deben ser localizadas y accedidas a través de múltiples proveedores, alineadas mediante proyecciones y resoluciones, y procesadas de manera eficiente. Luego, los resultados deben unirse en mapas geográficamente consistentes mientras la infraestructura se recupera de las fallas rutinarias de la computación distribuida.

Hoy, la plataforma puede ejecutar inferencia en áreas a escala continental en aproximadamente un día, procesando decenas de terabytes de imágenes a un costo de fracciones de centavo por kilómetro cuadrado. Desarrollarla significó enfrentar una serie de desafíos de ingeniería que otros que trabajan en sistemas geoespaciales a gran escala probablemente también encontrarán. Esta publicación recorre esos desafíos y las soluciones a las que llegamos.

Inferencia a escala continental por números: una cuadrícula de estadísticas que resume la ejecución de riesgo de incendios en Norteamérica
Inferencia a escala continental por números: una cuadrícula de estadísticas que resume la ejecución de riesgo de incendios en Norteamérica

Un mapa reciente de riesgo de incendios forestales generado en la plataforma OlmoEarth, con estadísticas.

La capacidad de procesar grandes volúmenes de datos permite obtener visualizaciones detalladas. La siguiente imagen muestra la resolución del análisis realizado.

Primer plano de un mapa de calor de riesgo de incendios sobre el terreno, desde azul (menor riesgo) hasta rojo (mayor riesgo)
Primer plano de un mapa de calor de riesgo de incendios sobre el terreno, desde azul (menor riesgo) hasta rojo (mayor riesgo)

¿Por qué la inferencia satelital es desafiante?

La mayoría de los modelos de ML reciben algunos megabytes de datos y producen un resultado en menos de un segundo; piense en LLMs procesando un párrafo de texto o modelos de visión artificial analizando una foto de un teléfono inteligente. La inferencia de observación terrestre opera a una escala muy diferente: un solo trabajo de ajuste fino de un modelo fundacional para obtener el máximo rendimiento puede mover terabytes de datos y ejecutarse durante horas. Las entradas pueden abarcar múltiples bandas espectrales, tipos de sensores y pasos de tiempo en una gran área geográfica. Pueden provenir de varios proveedores, cada uno utilizando diferentes proyecciones y resoluciones, y pueden incluir observaciones que faltan o están ocultas por nubes. El resultado es en sí mismo un mapa, por lo que cada predicción debe permanecer alineada precisamente con la misma proyección y cuadrícula de coordenadas que las áreas circundantes.

Incluso adquirir los datos puede ser un desafío importante. Los trabajos de predicción a menudo pasan más tiempo descargando y preparando imágenes que ejecutando el modelo en sí, lo que hace que los pipelines de datos eficientes sean críticos. Esos pipelines deben manejar I/O de alto volumen mientras proporcionan el cómputo necesario para reproyectar y remuestrear imágenes.

¿Cuál es el hardware adecuado para cada tarea?

Debido a que la adquisición y preparación de datos a menudo dominan el tiempo de ejecución de un trabajo de inferencia, asignar ese trabajo a GPUs dejaría al hardware más caro del sistema realizando tareas mejor adaptadas a CPUs. Por lo tanto, dividimos cada trabajo en tres etapas, cada una adaptada a un perfil de hardware distinto:

  • Adquisición y preprocesamiento de datos (CPU, alta I/O): Obtener, reproyectar, alinear y normalizar imágenes, luego escribirlas en un formato optimizado para una carga rápida durante la inferencia.
  • Inferencia (GPU): Ejecutar el paso forward del modelo y escribir resultados mínimamente procesados directamente en el almacenamiento.
  • Postprocesamiento (CPU): Unir los resultados por ventana, aplicar máscaras o reescalado, y exportarlos en formatos fáciles de usar como Zarr, GeoTIFF o GeoJSON.

La plataforma OlmoEarth distribuye estas etapas entre muchas máquinas mientras mantiene las GPUs totalmente utilizadas. Los cargadores de datos multiproceso alimentan continuamente a cada GPU, mientras que los resultados completados se transmiten directamente al almacenamiento de blobs.

Pipeline de inferencia de tres etapas como tarjetas: 01 Adquisición de datos en CPU, 02 Ejecución de inferencia en GPU, 03 Postprocesamiento en CPU
Pipeline de inferencia de tres etapas como tarjetas: 01 Adquisición de datos en CPU, 02 Ejecución de inferencia en GPU, 03 Postprocesamiento en CPU

El flujo de trabajo descrito asegura que cada componente del sistema trabaje a su máxima eficiencia según su arquitectura específica.

¿Cómo se gestionan cientos de trabajadores y miles de procesos?

OlmoEarth Run, la capa de ejecución de la plataforma para trabajos de inferencia a gran escala, divide la región geográfica cubierta por cada trabajo en particiones dimensionadas para instancias de cómputo individuales (workers), y luego subdivide esas particiones en ventanas más pequeñas que los modelos OlmoEarth procesan. Debido a que cada ventana puede ser manejada independientemente en un paso forward separado, el trabajo en una parte del mapa no necesita esperar a otra.

En la práctica, un área del tamaño de un estado podría convertirse en un centenar de particiones, mientras que una ejecución a escala continental puede convertirse en miles. Las particiones adyacentes se superponen ligeramente, y reconciliamos esa superposición cuando se ensamblan los resultados para que no aparezca ninguna costura en el ráster final.

Debido a que las particiones son independientes, la misma etapa puede ejecutarse en miles de instancias de cómputo a la vez. Recientemente utilizamos este enfoque para generar un mapa de riesgo de incendios forestales que cubre toda Norteamérica. En su punto máximo, la ejecución utilizó aproximadamente 19,600 CPUs y 994 GPUs en paralelo, con un rendimiento de red superior a 168 GB/s. Ese nivel de paralelismo redujo un estimado de 4,737 horas de cómputo serial a aproximadamente 30.5 horas de tiempo de reloj (wall-clock time); una aceleración de 155×.

Sin embargo, el fan-out no es ilimitado. Más trabajadores presionan contra las cuotas de la nube, por lo que el paralelismo es una perilla por ejecución, una de varias que podemos ajustar en trabajos individuales. La resolución de salida intercambia volumen de datos y cómputo por detalle; el tamaño del modelo intercambia tiempo de GPU por precisión; el almacenamiento en caché de imágenes crudas intercambia almacenamiento por velocidad en ejecuciones repetidas sobre la misma área. La configuración correcta depende de la tarea —y del presupuesto— disponible.

Particionamiento de dos etapas de OlmoEarth: la etapa 1 divide una solicitud en particiones por máquina, la etapa 2 subdivide cada celda en ventanas de modelo
Particionamiento de dos etapas de OlmoEarth: la etapa 1 divide una solicitud en particiones por máquina, la etapa 2 subdivide cada celda en ventanas de modelo

La arquitectura de dos etapas permite escalar el proceso de forma modular según las necesidades de cada proyecto específico.

¿Cómo se obtienen los píxeles correctos?

Dada una región geográfica y un rango de tiempo, la plataforma primero debe determinar qué escenas satelitales deben alimentar al modelo. Eso significa identificar qué se capturó, dónde y cuándo a través de proveedores con diferentes catálogos, formatos y retrasos de publicación. Los criterios de selección también dependen de la fuente: para imágenes ópticas como Sentinel-2, generalmente queremos las escenas menos nubladas disponibles, mientras que para el radar de apertura sintética, los canales de polarización disponibles.

Vía Hugging Face.