GeoSeg 框架解析
总的来看：
config 配置文件
geoseg 定义了核心组件：models datasets losses
tools 辅助工具箱
config :
里面包含多个遥感数据集，内部包含不同的模型
看一个例子
导入
from torch.utils.data import DataLoader
from geoseg.losses import *
from geoseg.datasets.vaihingen_dataset import *
from geoseg.models.UNetFormer import UNetFormer
from tools.utils import Lookahead, process_model_params
训练超参数
from torch.utils.data import DataLoader
from geoseg.losses import *
from geoseg.datasets.vaihingen_dataset import *
from geoseg.models.UNetFormer import UNetFormer
from tools.utils import Lookahead, process_model_params
网络定义
net = UNetFormer(num_classes=num_classes)
损失函数
loss = UnetFormerLoss(ignore_index=ignore_index)
use_aux_loss = True
数据加载器
train_dataset = VaihingenDataset(data_root='data/vaihingen/train', ...)
val_dataset = VaihingenDataset(...)
...
train_loader = DataLoader(dataset=train_dataset, ...)
val_loader = DataLoader(dataset=val_dataset, ...)
优化器
layerwise_params = {"backbone.*": dict(lr=backbone_lr, ...)}
net_params = process_model_params(net, ...)
base_optimizer = torch.optim.AdamW(net_params, lr=lr, ...)
optimizer = Lookahead(base_optimizer)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(...)
geoseg
models—>所有神经网络骨架库
datasets–>数据集库
losses —>损失函数库
datasets
除了每个数据集的datasets 之外，有个transform是数据增强策略
transform: 各种各样的数据增强策略
看一个dataset例子：
全局常量
类别
CLASSES = ('ImSurf', 'Building', 'LowVeg', 'Tree', 'Car', 'Clutter')
可视化颜色
PALETTE = [[255, 255, 255], [0, 0, 255], [0, 255, 255], [0, 255, 0], [255, 204, 0], [255, 0, 0]]ORIGIN_IMG_SIZE = (1024, 1024)
INPUT_IMG_SIZE = (1024, 1024)
TEST_IMG_SIZE = (1024, 1024)
数据增强流水线
def get_training_transform(): ...
def train_aug(img, mask): ...
def get_val_transform(): ...
def val_aug(img, mask): ...
这些都是对数据进行数据增强
核心数据集类
class VaihingenDataset(Dataset):def __init__(...): ...构造函数，保存传入参数，调用get_img_ids读取images_1024目录下的所有文件名，并且保存在self.img_ids 当中def __getitem__(...): ...def __len__(...): ...def get_img_ids(...): ...用来确保img_dir 和mask_dir能对应的上def load_img_and_mask(...): ...加载数据def load_mosaic_img_and_mask(...): ...高级数据增强可视化工具
def show_img_mask_seg(...): ...
def show_seg(...): ...
def show_mask(...): ...
losses
各种各样loss 函数
models：定义了不同的models 神经网络
看一个例子
这个例子没有encoder 模块是因为，他的encoder 模块就只是backbone 特征提取，如果特征提取后加入混合注意力之类的东西，就应该写一个encoder类
此外，models 一般都是把组件，基础架构模块定义好，核心模块定义好，在encoder 和 decoder 内部，按照论文结构排列组合
encoder 和 decoder 会在最后的模型封装中，在forward 中应用
基础架构模块
class ConvBNReLU(nn.Sequential): ...
class ConvBN(nn.Sequential): ...
class Conv(nn.Sequential): ...
class SeparableConvBNReLU(nn.Sequential): ...
class Mlp(nn.Module): ...核心注意力模块
class GlobalLocalAttention(nn.Module): ...
Transformer 模块
class GlobalLocalAttention(nn.Module): ...
解码辅助模块
class WF(nn.Module): ...
class FeatureRefinementHead(nn.Module): ...
class AuxHead(nn.Module): ...
解码器
class Decoder(nn.Module): ...
最终模型
class UNetFormer(nn.Module): ...
构造函数和前向传播tools
数据预处理脚本
模型配置脚本
模型评估脚本