针对地球观测数据优化云到 GPU 深度学习吞吐量
计算机视觉与模式识别
2025-06-09 v1
摘要
在 petabyte 规模的地球观测 (EO) 数据上训练深度学习模型需要将计算资源与数据存储分离。然而，标准的 PyTorch 数据加载器无法在从云存储直接流式传输 GeoTIFF 文件时保持现代 GPU 的利用率。在本工作中，我们对从云对象存储和本地 SSD 加载 GeoTIFF 的吞吐量进行基准测试，系统性地测试不同的加载器配置和数据参数。我们专注于与瓷砖对齐的读取和工作线程池，使用贝叶斯优化为每种存储类型找到最佳设置。我们的优化配置将远程数据加载吞吐量提高了 20 倍，将本地吞吐量提高了 4 倍。在三个公共 EO 基准数据集上，使用优化后的远程加载训练的模型，在相同的预算时间内实现与本地训练相同的准确率。我们将验证 IoU 提高了 6-15%，GPU 利用率维持在 85-95%，而标准配置下仅为 0-30%。代码已公开可用，地址为 https://github.com/microsoft/pytorch-cloud-geotiff-optimization
引用
@article{arxiv.2506.06235,
title = {Optimizing Cloud-to-GPU Throughput for Deep Learning With Earth Observation Data},
author = {Akram Zaytar and Caleb Robinson and Girmaw Abebe Tadesse and Tammy Glazer and Gilles Hacheme and Anthony Ortiz and Rahul M Dodhia and Juan M Lavista Ferres},
journal= {arXiv preprint arXiv:2506.06235},
year = {2025}
}