GeoTorchAI: 用于大规模时空深度学习模型训练与使用的框架

晏彤钰Mighty

于 2025-04-05 10:30:35 发布

阅读量727

点赞数 8

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_01181/article/details/147008360

GeoTorchAI: 用于大规模时空深度学习模型训练与使用的框架

GeoTorchAI GeoTorchAI: A Framework for Training and Using Spatiotemporal Deep Learning Models at Scale 项目地址: https://gitcode.com/gh_mirrors/ge/GeoTorchAI

1. 项目介绍

GeoTorchAI 是一个基于 PyTorch 和 Apache Sedona 的时空深度学习框架。它使得时空机器学习实践者能够轻松高效地实现针对栅格影像数据集和时空非影像数据集的深度学习模型。针对栅格影像数据集的深度学习应用包括卫星影像分类和卫星图像分割。在时空非影像数据集上的深度学习应用主要是预测任务，包括但不限于交通量预测、出租车/自行车流量预测、降水预测和天气预报。

2. 项目快速启动

首先，确保已经安装了 Python 和 pip。然后，可以通过以下命令安装 GeoTorchAI：

pip install geotorchai

快速启动示例：卫星影像分类

以下是一个使用 GeoTorchAI 进行卫星影像分类的简单示例：

# 导入必要的库
import torch
import numpy as np
from geotorchai.datasets import EuroSAT
from geotorchai.models import DeepSatV2
from torch.utils.data import DataLoader, SubsetRandomSampler
from torch.optim import Adam
from torch.nn import CrossEntropyLoss

# 加载 EuroSAT 数据集
full_data = EuroSAT(root="data/eurosat", download=True, include_additional_features=True)

# 划分训练和验证数据集
dataset_size = len(full_data)
indices = list(range(dataset_size))
split = int(np.floor(0.2 * dataset_size))
np.random.seed(42)
np.random.shuffle(indices)
train_indices, val_indices = indices[split:], indices[:split]
train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_indices)
valid_sampler = torch.utils.data.sampler.SubsetRandomSampler(val_indices)
train_loader = DataLoader(full_data, batch_size=16, sampler=train_sampler)
val_loader = DataLoader(full_data, batch_size=16, sampler=valid_sampler)

# 初始化设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 初始化模型、损失函数和优化器
model = DeepSatV2(in_channels=13, in_height=64, in_width=64, num_classes=10, num_filtered_features=len(full_data.ADDITIONAL_FEATURES))
loss_fn = CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=0.0002)

# 将模型和损失函数移动到设备上
model.to(device)
loss_fn.to(device)

# 训练模型
for epoch in range(1):
    model.train()
    for inputs, labels, features in train_loader:
        inputs, features = inputs.to(device), features.type(torch.FloatTensor).to(device)
        labels = labels.to(device)
        outputs = model(inputs, features)
        loss = loss_fn(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# 验证模型
model.eval()
total_sample = 0
correct = 0
with torch.no_grad():
    for inputs, labels, features in val_loader:
        inputs, features = inputs.to(device), features.type(torch.FloatTensor).to(device)
        labels = labels.to(device)
        outputs = model(inputs, features)
        _, predicted = outputs.max(1)
        total_sample += len(labels)
        correct += predicted.eq(labels).sum().item()

val_accuracy = 100 * correct / total_sample
print(f"Validation Accuracy: {val_accuracy}%")