TorchSpatial: A Python Package for Spatial Representation Learning and Geo-Aware Model Development

Qian Cao, Nemin Wu, Zhangyu Wang, Zeping Liu, Yanlin Qi, Jielu Zhang, Joshua Ni, Xiaobai Yao, Hongxu Ma, Lan Mu, Stefano Ermon, Tanuja Ganu, Akshay Nambi, Ni Lao, Gengchen Mai

Abstract: Spatial representation learning (SRL) focuses on developing spatial embeddings from various forms of spatial data, such as points, polylines, polygons, graphs, networks, and images without any additional feature engineering or data conversion step. Effective spatial representation is fundamental for a wide range of downstream geospatial applications, including species distribution modeling, satellite image classification, point cloud classification and segementation, trajectory synthesis, building footprint extraction, and cartographic generalization. Despite the widespread use of SRL as a cornerstone for many spatially-aware AI models, there is still no comprehensive package shared across the community that provides ready-made code to support the implementation and reproduction of SRL model development. To fill this void, we present TorchSpatial, a Python package designed to support the encoding of spatial data, starting with location (point) encoding, a fundamental data type in SRL. TorchSpatial includes two key components: 1) We present TorchSpatial, an SRL framework supporting the development of location encoders. TorchSpatial now integrates 15 widely-used encoders and essential encoder components, ensuring scalability and reproducibility for future developments; 2) We establish a ready-to-use workflow that takes the input hyperparameters and outputs the model inference results and evaluation across geo-aware image classification and regression tasks with access to 17 datasets. We believe TorchSpatial will foster future advancement of SRL and spatial fairness in GeoAI research. The TorchSpatial SRL framework and inference models are available at https://github.com/seai-lab/TorchSpatial.