【深度学习|地学应用】使用SimCLR对遥感数据（如AID）训练+评估完整代码流程，附代码。

985在读小水博一枚呀~

墨衍会员 · AI 创作全网分发码龄2年

关注

原创最新推荐文章于 2026-05-14 15:05:45 发布 928 阅读 AI 写同款· GEO 优化 ›

论文阅读——GeoChat（cvpr2024）阅读详情

【深度学习|地学应用】使用SimCLR（a simple framework for contrastive learning of visual representations）对遥感数据（如AID）训练+评估完整代码流程，附代码。

文章目录

【深度学习|地学应用】使用SimCLR（a simple framework for contrastive learning of visual representations）对遥感数据（如AID）训练+评估完整代码流程，附代码。
前言

欢迎铁子们点赞、关注、收藏！
祝大家逢考必过！逢投必中！上岸上岸上岸！upupup

大多数高校硕博生毕业要求需要参加学术会议，发表EI或者SCI检索的学术论文会议论文。详细信息可关注VX “学术会议小灵通”或参考学术信息专栏：https://blog.csdn.net/2401_89898861/article/details/147776758

前言

下面是一个使用 SimCLR 方法对 AID 数据集进行自监督预训练与下游分类微调的完整流程，包括环境配置、数据准备、模型构建、训练与评估。

该代码基于 PyTorch 编写，适用于遥感图像自监督学习实验。

✅ 一、环境准备

pip install torch torchvision scikit-learn tqdm

✅ 二、数据准备（AID 数据集）

你需要先下载 AID 遥感图像分类数据集，每个类有220张图像，总共30个类别（彩色图像，600×600）。
我们假设你将其解压到了 ./AID 文件夹下：

./AID/
    ├── airport/
    ├── bareland/
    ├── ...

✅ 三、数据加载与增强

from torchvision import transforms, datasets
from torch.utils.data import Dataset
import random

class SimCLRDataset(Dataset):
    def __init__(self, root, transform):
        self.dataset = datasets.ImageFolder(root=root)
        self.transform = transform

    def __getitem__(self, idx):
        img, _ = self.dataset[idx]
        xi = self.transform(img)
        xj = self.transform(img)
        return xi, xj

    def __len__(self):
        return len(self.dataset)

# SimCLR增强：随机裁剪 + 颜色扰动
simclr_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.RandomApply([transforms.ColorJitter(0.4,0.4,0.4,0.1)], p=0.8),
    transforms.RandomGrayscale(p=0.2),
    transforms.ToTensor(),
])

train_dataset = SimCLRDataset("./AID", simclr_transform)

✅ 四、SimCLR 模型构建

import torch.nn as nn
import torchvision.models as models

class SimCLR(nn.Module):
    def __init__(self, base_model='resnet18', projection_dim=128):
        super().__init__()
        self.encoder = models.__dict__[base_model](pretrained=False)
        num_ftrs = self.encoder.fc.in_features
        self.encoder.fc = nn.Identity()

        self.projector = nn.Sequential(
            nn.Linear(num_ftrs, 256),
            nn.ReLU(),
            nn.Linear(256, projection_dim)
        )

    def forward(self, x):
        h = self.encoder(x)
        z = self.projector(h)
        return z

✅ 五、NT-Xent 对比损失函数

import torch
import torch.nn.functional as F

def nt_xent_loss(z_i, z_j, temperature=0.5):
    N = z_i.shape[0]
    z = torch.cat([z_i, z_j], dim=0)
    z = F.normalize(z, dim=1)

    sim = torch.mm(z, z.T)  # [2N, 2N]
    sim /= temperature

    labels = torch.arange(N).repeat(2)
    labels = (labels.unsqueeze(0) == labels.unsqueeze(1)).float().cuda()

    mask = torch.eye(2 * N, dtype=torch.bool).cuda()
    labels = labels[~mask].view(2 * N, -1)
    sim = sim[~mask].view(2 * N, -1)

    positives = sim[labels.bool()].view(2 * N, -1)
    negatives = sim[~labels.bool()].view(2 * N, -1)

    logits = torch.cat([positives, negatives], dim=1)
    labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()

    return F.cross_entropy(logits, labels)

✅ 六、预训练 SimCLR 模型

from torch.utils.data import DataLoader
import torch.optim as optim
from tqdm import tqdm

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimCLR().to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)

train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)

for epoch in range(100):
    model.train()
    total_loss = 0
    for xi, xj in tqdm(train_loader):
        xi, xj = xi.to(device), xj.to(device)
        zi = model(xi)
        zj = model(xj)
        loss = nt_xent_loss(zi, zj)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"[Epoch {epoch}] Loss: {total_loss / len(train_loader):.4f}")

✅ 七、下游分类微调

提取冻结后的 encoder 特征，训练一个线性分类器（linear probe）：

class LinearClassifier(nn.Module):
    def __init__(self, encoder, num_classes=30):
        super().__init__()
        self.encoder = encoder
        for param in self.encoder.parameters():
            param.requires_grad = False  # 冻结
        self.fc = nn.Linear(512, num_classes)

    def forward(self, x):
        x = self.encoder(x)
        return self.fc(x)

用标准 ImageFolder 加载有标签的训练数据：

supervised_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])
train_sup = datasets.ImageFolder("./AID", transform=supervised_transform)
train_loader = DataLoader(train_sup, batch_size=64, shuffle=True)

训练线性分类器：

model = LinearClassifier(model.encoder).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=1e-3)

for epoch in range(20):
    model.train()
    total = 0
    correct = 0
    for imgs, labels in train_loader:
        imgs, labels = imgs.to(device), labels.to(device)
        outputs = model(imgs)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total += labels.size(0)
        correct += (outputs.argmax(1) == labels).sum().item()
    print(f"Epoch {epoch}: Acc = {correct / total * 100:.2f}%")

✅ 八、总结

模块	描述
数据增强	生成双视图用于对比学习
模型架构	ResNet + MLP投影头
预训练目标	最大化相同图像增强视图的相似度
下游任务	分类（冻结 encoder + 线性分类器）