PyTorch遥感深度学习实战：从CNN原理到目标检测与图像分割

在遥感影像分析领域，深度学习技术正发挥着越来越重要的作用。无论是城市规划、农业监测还是灾害评估，都需要从海量遥感数据中快速准确地提取有价值信息。本文将基于PyTorch框架，完整介绍遥感深度学习的全流程实战，涵盖CNN原理、目标检测、图像分割和优化技巧四大核心模块，并提供可直接运行的完整代码和数据集处理方案。

1. 遥感深度学习基础与环境搭建

1.1 遥感影像特点与深度学习应用场景

遥感影像与传统图像相比具有显著差异：多光谱/高光谱通道、大尺寸、空间分辨率各异、地物尺度变化大等特性。这些特点决定了遥感深度学习需要特殊的处理方法和模型架构。

典型应用场景包括：

• 目标检测：建筑物、车辆、船舶等人工地物的自动识别
• 语义分割：土地覆盖分类、道路提取、植被监测
• 变化检测：城市扩张、灾害损毁评估
• 超分辨率重建：提升低分辨率影像质量

1.2 PyTorch环境配置与依赖安装

推荐使用Anaconda创建独立的Python环境，确保依赖包版本兼容性：

# 创建conda环境 conda create -n rsdl python=3.8 conda activate rsdl # 安装PyTorch（根据CUDA版本选择） pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html # 安装遥感处理相关库 pip install opencv-python rasterio gdal pillow pip install scikit-learn scikit-image matplotlib seaborn pip install albumentations tqdm tensorboard

1.3 数据集准备与预处理

遥感数据集通常需要特殊处理，以下是一个通用的数据加载类：

import torch from torch.utils.data import Dataset import rasterio import numpy as np import albumentations as A from albumentations.pytorch import ToTensorV2 class RemoteSensingDataset(Dataset): def __init__(self, image_paths, mask_paths=None, transform=None): self.image_paths = image_paths self.mask_paths = mask_paths self.transform = transform def __len__(self): return len(self.image_paths) def __getitem__(self, idx): # 读取遥感影像（支持多光谱） with rasterio.open(self.image_paths[idx]) as img: image = img.read() # [C, H, W] image = np.transpose(image, (1, 2, 0)) # [H, W, C] # 数据增强和预处理 if self.transform: augmented = self.transform(image=image) image = augmented['image'] # 如果有标注数据 if self.mask_paths is not None: with rasterio.open(self.mask_paths[idx]) as mask: mask_data = mask.read(1) # 单通道标注 if self.transform: augmented = self.transform(image=image, mask=mask_data) image, mask_data = augmented['image'], augmented['mask'] return image, mask_data.long() return image # 数据增强配置 def get_train_transform(): return A.Compose([ A.RandomRotate90(p=0.5), A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5), A.RandomBrightnessContrast(p=0.2), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2(), ])

2. CNN原理与遥感特征提取

2.1 卷积神经网络基础架构

CNN通过局部连接和权值共享有效提取空间特征，特别适合遥感影像处理。以下是基础CNN模型实现：

import torch.nn as nn import torch.nn.functional as F class BasicCNN(nn.Module): def __init__(self, in_channels=3, num_classes=10): super(BasicCNN, self).__init__() # 特征提取层 self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1) self.bn1 = nn.BatchNorm2d(64) self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1) self.bn2 = nn.BatchNorm2d(128) self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1) self.bn3 = nn.BatchNorm2d(256) # 全局平均池化替代全连接层 self.global_avg_pool = nn.AdaptiveAvgPool2d((1, 1)) self.classifier = nn.Linear(256, num_classes) # Dropout防止过拟合 self.dropout = nn.Dropout(0.5) def forward(self, x): # 第一层卷积块 x = F.relu(self.bn1(self.conv1(x))) x = F.max_pool2d(x, 2) # 第二层卷积块 x = F.relu(self.bn2(self.conv2(x))) x = F.max_pool2d(x, 2) # 第三层卷积块 x = F.relu(self.bn3(self.conv3(x))) x = F.max_pool2d(x, 2) # 全局特征提取 x = self.global_avg_pool(x) x = x.view(x.size(0), -1) x = self.dropout(x) x = self.classifier(x) return x

2.2 遥感影像特征提取技巧

遥感影像特征提取需要考虑空间尺度和光谱特性：

class MultiScaleCNN(nn.Module): """多尺度特征提取网络，适合不同大小的遥感目标""" def __init__(self, in_channels=3): super(MultiScaleCNN, self).__init__() # 小尺度特征提取（细节特征） self.small_scale = nn.Sequential( nn.Conv2d(in_channels, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.Conv2d(64, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), ) # 中尺度特征提取 self.medium_scale = nn.Sequential( nn.Conv2d(in_channels, 64, 5, padding=2), nn.BatchNorm2d(64), nn.ReLU(), nn.Conv2d(64, 64, 5, padding=2), nn.BatchNorm2d(64), nn.ReLU(), ) # 大尺度特征提取（上下文特征） self.large_scale = nn.Sequential( nn.Conv2d(in_channels, 64, 7, padding=3), nn.BatchNorm2d(64), nn.ReLU(), nn.Conv2d(64, 64, 7, padding=3), nn.BatchNorm2d(64), nn.ReLU(), ) # 特征融合 self.feature_fusion = nn.Sequential( nn.Conv2d(192, 256, 3, padding=1), nn.BatchNorm2d(256), nn.ReLU(), nn.Dropout2d(0.3) ) def forward(self, x): small_feat = self.small_scale(x) medium_feat = self.medium_scale(x) large_feat = self.large_scale(x) # 多尺度特征拼接 fused_feat = torch.cat([small_feat, medium_feat, large_feat], dim=1) output = self.feature_fusion(fused_feat) return output

2.3 迁移学习在遥感中的应用

使用预训练模型可以显著提升小样本遥感数据的训练效果：

import torchvision.models as models class TransferLearningModel(nn.Module): def __init__(self, num_classes, pretrained=True): super(TransferLearningModel, self).__init__() # 使用ResNet预训练 backbone self.backbone = models.resnet50(pretrained=pretrained) # 替换第一层卷积，适应多光谱输入 original_conv1 = self.backbone.conv1 self.backbone.conv1 = nn.Conv2d( 3, 64, kernel_size=7, stride=2, padding=3, bias=False ) # 复制预训练权重（仅RGB通道） if pretrained: with torch.no_grad(): self.backbone.conv1.weight[:, :3] = original_conv1.weight # 多光谱通道使用RGB均值初始化 if 3 > 3: self.backbone.conv1.weight[:, 3:] = original_conv1.weight.mean(dim=1, keepdim=True) # 修改分类器 in_features = self.backbone.fc.in_features self.backbone.fc = nn.Sequential( nn.Dropout(0.5), nn.Linear(in_features, 512), nn.ReLU(), nn.Dropout(0.3), nn.Linear(512, num_classes) ) def forward(self, x): return self.backbone(x)

3. 遥感目标检测实战

3.1 两阶段目标检测器实现

基于Faster R-CNN框架实现遥感目标检测：

import torchvision from torchvision.models.detection import FasterRCNN from torchvision.models.detection.rpn import AnchorGenerator class RemoteSensingDetector: def __init__(self, num_classes, backbone='resnet50'): self.num_classes = num_classes self.backbone = backbone self.model = self._build_model() def _build_model(self): # 加载预训练backbone if self.backbone == 'resnet50': backbone = torchvision.models.resnet50(pretrained=True) backbone = nn.Sequential(*list(backbone.children())[:-2]) # 定义RPN锚点生成器（适应遥感目标尺度） anchor_sizes = ((32,), (64,), (128,), (256,), (512,)) aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes) anchor_generator = AnchorGenerator( sizes=anchor_sizes, aspect_ratios=aspect_ratios ) # ROI对齐层 roi_pooler = torchvision.ops.MultiScaleRoIAlign( featmap_names=['0'], output_size=7, sampling_ratio=2 ) # 构建Faster R-CNN模型 model = FasterRCNN( backbone, num_classes=self.num_classes, rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler ) return model def train_model(self, train_loader, val_loader, num_epochs=50): optimizer = torch.optim.SGD( self.model.parameters(), lr=0.005, momentum=0.9, weight_decay=0.0005 ) lr_scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=10, gamma=0.1 ) self.model.train() for epoch in range(num_epochs): total_loss = 0 for images, targets in train_loader: optimizer.zero_grad() loss_dict = self.model(images, targets) losses = sum(loss for loss in loss_dict.values()) losses.backward() optimizer.step() total_loss += losses.item() lr_scheduler.step() print(f'Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/len(train_loader):.4f}')

3.2 单阶段目标检测器优化

YOLO系列在遥感检测中具有速度优势，以下是简化版实现：

class YOLOv5Like(nn.Module): """基于YOLOv5架构的遥感目标检测器""" def __init__(self, num_classes, anchors=None): super(YOLOv5Like, self).__init__() self.num_classes = num_classes # 默认锚点框（根据遥感目标调整） if anchors is None: self.anchors = torch.tensor([ [10, 13], [16, 30], [33, 23], # 小目标 [30, 61], [62, 45], [59, 119], # 中目标 [116, 90], [156, 198], [373, 326] # 大目标 ]) / 640.0 # 归一化 # Backbone: CSPDarknet简化版 self.backbone = self._build_backbone() # Neck: PANet特征金字塔 self.neck = self._build_neck() # Head: 检测头 self.head = self._build_head() def _build_backbone(self): return nn.Sequential( # 初始卷积块 nn.Conv2d(3, 32, 3, padding=1), nn.BatchNorm2d(32), nn.SiLU(), nn.MaxPool2d(2), # 多个CSP块 self._csp_block(32, 64, 1), nn.MaxPool2d(2), self._csp_block(64, 128, 3), nn.MaxPool2d(2), self._csp_block(128, 256, 3), nn.MaxPool2d(2), self._csp_block(256, 512, 1), ) def _csp_block(self, in_c, out_c, n): """跨阶段部分网络块""" layers = [] # 实现细节... return nn.Sequential(*layers)

3.3 目标检测数据预处理与增强

遥感目标检测需要特殊的数据增强策略：

class DetectionTransform: """目标检测专用数据增强""" def __init__(self, image_size=640, is_training=True): self.image_size = image_size self.is_training = is_training def __call__(self, image, targets): if self.is_training: # 训练时的增强策略 transform = A.Compose([ A.Resize(self.image_size, self.image_size), A.HorizontalFlip(p=0.5), A.RandomBrightnessContrast(p=0.2), A.ShiftScaleRotate( shift_limit=0.1, scale_limit=0.1, rotate_limit=15, p=0.5 ), A.OneOf([ A.GaussNoise(var_limit=(10.0, 50.0)), A.Blur(blur_limit=3), ], p=0.3), ], bbox_params=A.BboxParams( format='pascal_voc', label_fields=['labels'] )) else: # 验证/测试时的处理 transform = A.Compose([ A.Resize(self.image_size, self.image_size), ], bbox_params=A.BboxParams( format='pascal_voc', label_fields=['labels'] )) transformed = transform(image=image, bboxes=targets['boxes'], labels=targets['labels']) return transformed['image'], { 'boxes': torch.tensor(transformed['bboxes'], dtype=torch.float32), 'labels': torch.tensor(transformed['labels'], dtype=torch.long) }

4. 遥感图像分割技术

4.1 U-Net架构与遥感适配

U-Net在遥感分割中表现优异，以下是完整实现：

class UNet(nn.Module): """适用于遥感图像分割的U-Net架构""" def __init__(self, in_channels=3, out_channels=1, features=[64, 128, 256, 512]): super(UNet, self).__init__() self.encoder = nn.ModuleList() self.decoder = nn.ModuleList() self.pool = nn.MaxPool2d(kernel_size=2, stride=2) # Encoder路径 for feature in features: self.encoder.append(UNet._block(in_channels, feature)) in_channels = feature # Bottleneck self.bottleneck = UNet._block(features[-1], features[-1] * 2) # Decoder路径 for feature in reversed(features): self.decoder.append( nn.ConvTranspose2d(feature * 2, feature, kernel_size=2, stride=2) ) self.decoder.append(UNet._block(feature * 2, feature)) # 最终卷积层 self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1) @staticmethod def _block(in_channels, features): return nn.Sequential( nn.Conv2d(in_channels, features, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(features), nn.ReLU(inplace=True), nn.Conv2d(features, features, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(features), nn.ReLU(inplace=True), ) def forward(self, x): skip_connections = [] # Encoder for encode in self.encoder: x = encode(x) skip_connections.append(x) x = self.pool(x) # Bottleneck x = self.bottleneck(x) skip_connections = skip_connections[::-1] # Decoder for idx in range(0, len(self.decoder), 2): x = self.decoder[idx](x) skip_connection = skip_connections[idx//2] # 跳跃连接（处理尺寸不匹配） if x.shape != skip_connection.shape: x = F.interpolate(x, size=skip_connection.shape[2:], mode='bilinear') concat_skip = torch.cat((skip_connection, x), dim=1) x = self.decoder[idx+1](concat_skip) return torch.sigmoid(self.final_conv(x))

4.2 DeepLabv3+在遥感中的应用

DeepLabv3+通过空洞卷积捕获多尺度上下文信息：

class DeepLabV3Plus(nn.Module): """DeepLabv3+遥感图像分割模型""" def __init__(self, num_classes, backbone='resnet50', output_stride=16): super(DeepLabV3Plus, self).__init__() self.backbone = self._build_backbone(backbone, output_stride) self.aspp = ASPP(2048, 256, output_stride) self.decoder = Decoder(num_classes, 256) def _build_backbone(self, backbone_name, output_stride): # 基于ResNet构建backbone，支持空洞卷积 if backbone_name == 'resnet50': backbone = models.resnet50(pretrained=True) # 修改最后两个block为空洞卷积 if output_stride == 16: backbone.layer4[0].conv1.stride = (1, 1) backbone.layer4[0].downsample[0].stride = (1, 1) return backbone def forward(self, x): # 提取低级特征 low_level_feat = self.backbone.layer1(x) low_level_feat = self.backbone.layer2(low_level_feat) # 高级特征提取 x = self.backbone.layer3(low_level_feat) x = self.backbone.layer4(x) # ASPP模块 x = self.aspp(x) # 解码器 x = self.decoder(x, low_level_feat) # 上采样到原图尺寸 x = F.interpolate(x, scale_factor=4, mode='bilinear', align_corners=True) return x class ASPP(nn.Module): """空洞空间金字塔池化""" def __init__(self, in_channels, out_channels, output_stride): super(ASPP, self).__init__() if output_stride == 16: dilations = [1, 6, 12, 18] elif output_stride == 8: dilations = [1, 12, 24, 36] self.aspp_blocks = nn.ModuleList([ nn.Sequential( nn.Conv2d(in_channels, out_channels, 1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU() ) ]) for dilation in dilations[1:]: self.aspp_blocks.append( nn.Sequential( nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU() ) ) self.global_avg_pool = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_channels, out_channels, 1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU() ) self.conv1 = nn.Conv2d(out_channels * 5, out_channels, 1, bias=False) self.bn = nn.BatchNorm2d(out_channels) self.relu = nn.ReLU() self.dropout = nn.Dropout(0.5) def forward(self, x): # 各分支处理 aspp_outs = [] for aspp_block in self.aspp_blocks: aspp_outs.append(aspp_block(x)) # 全局平均池化分支 global_feat = self.global_avg_pool(x) global_feat = F.interpolate(global_feat, size=x.shape[2:], mode='bilinear', align_corners=True) aspp_outs.append(global_feat) # 特征拼接和融合 x = torch.cat(aspp_outs, dim=1) x = self.conv1(x) x = self.bn(x) x = self.relu(x) x = self.dropout(x) return x

4.3 分割模型训练与评估

完整的训练流程和评估指标：

class SegmentationTrainer: def __init__(self, model, device, criterion, optimizer): self.model = model.to(device) self.device = device self.criterion = criterion self.optimizer = optimizer self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', patience=5, factor=0.5 ) def train_epoch(self, dataloader): self.model.train() running_loss = 0.0 iou_scores = [] for batch_idx, (images, masks) in enumerate(dataloader): images = images.to(self.device) masks = masks.to(self.device) self.optimizer.zero_grad() outputs = self.model(images) loss = self.criterion(outputs, masks) loss.backward() self.optimizer.step() running_loss += loss.item() # 计算IoU preds = torch.argmax(outputs, dim=1) iou = self.calculate_iou(preds, masks) iou_scores.append(iou) if batch_idx % 10 == 0: print(f'Batch {batch_idx}, Loss: {loss.item():.4f}, IoU: {iou:.4f}') avg_loss = running_loss / len(dataloader) avg_iou = np.mean(iou_scores) return avg_loss, avg_iou def calculate_iou(self, preds, targets): """计算交并比""" intersection = (preds & targets).float().sum((1, 2)) union = (preds | targets).float().sum((1, 2)) iou = (intersection + 1e-6) / (union + 1e-6) return iou.mean().item() # 多类别分割损失函数 class MultiClassDiceLoss(nn.Module): def __init__(self, weight=None, smooth=1e-6): super(MultiClassDiceLoss, self).__init__() self.smooth = smooth self.weight = weight def forward(self, pred, target): pred = F.softmax(pred, dim=1) target_onehot = F.one_hot(target, pred.size(1)).permute(0, 3, 1, 2) intersection = (pred * target_onehot).sum(dim=(2, 3)) union = pred.sum(dim=(2, 3)) + target_onehot.sum(dim=(2, 3)) dice = (2. * intersection + self.smooth) / (union + self.smooth) if self.weight is not None: dice = dice * self.weight return 1 - dice.mean()

5. 模型优化技巧与实战策略

5.1 学习率调度与优化器选择

针对遥感数据特点的优化策略：

def create_optimizer(model, optimizer_type='adamw', lr=1e-4, weight_decay=1e-4): """创建适合遥感任务的优化器""" if optimizer_type == 'adamw': optimizer = torch.optim.AdamW( model.parameters(), lr=lr, weight_decay=weight_decay, betas=(0.9, 0.999) ) elif optimizer_type == 'sgd': optimizer = torch.optim.SGD( model.parameters(), lr=lr, momentum=0.9, weight_decay=weight_decay, nesterov=True ) return optimizer def create_scheduler(optimizer, scheduler_type='cosine', epochs=100): """学习率调度器""" if scheduler_type == 'cosine': scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=epochs ) elif scheduler_type == 'step': scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=30, gamma=0.1 ) elif scheduler_type == 'plateau': scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', patience=10, factor=0.5 ) return scheduler # 组合优化策略 class OptimizerManager: def __init__(self, model, init_lr=1e-3): self.model = model self.optimizer = create_optimizer(model, 'adamw', init_lr) self.scheduler = create_scheduler(self.optimizer, 'cosine') def step(self, loss=None): self.optimizer.step() if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): self.scheduler.step(loss) else: self.scheduler.step() def zero_grad(self): self.optimizer.zero_grad()

5.2 过拟合防止与正则化技术

遥感数据量通常有限，需要有效的正则化：

class RegularizationTechniques: """正则化技术集合""" @staticmethod def label_smoothing(pred, target, smoothing=0.1): """标签平滑""" n_classes = pred.size(-1) one_hot = torch.full_like(pred, smoothing / (n_classes - 1)) one_hot.scatter_(1, target.unsqueeze(1), 1.0 - smoothing) return F.kl_div(F.log_softmax(pred, dim=1), one_hot, reduction='batchmean') @staticmethod def cutmix_data(x, y, alpha=1.0): """CutMix数据增强""" lam = np.random.beta(alpha, alpha) batch_size = x.size()[0] index = torch.randperm(batch_size) y_a, y_b = y, y[index] bbx1, bby1, bbx2, bby2 = rand_bbox(x.size(), lam) x[:, :, bbx1:bbx2, bby1:bby2] = x[index, :, bbx1:bbx2, bby1:bby2] lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (x.size()[-1] * x.size()[-2])) return x, y_a, y_b, lam @staticmethod def mixup_data(x, y, alpha=1.0): """MixUp数据增强""" if alpha > 0: lam = np.random.beta(alpha, alpha) else: lam = 1 batch_size = x.size()[0] index = torch.randperm(batch_size) mixed_x = lam * x + (1 - lam) * x[index, :] y_a, y_b = y, y[index] return mixed_x, y_a, y_b, lam def rand_bbox(size, lam): """生成随机边界框""" W = size[2] H = size[3] cut_rat = np.sqrt(1. - lam) cut_w = int(W * cut_rat) cut_h = int(H * cut_rat) cx = np.random.randint(W) cy = np.random.randint(H) bbx1 = np.clip(cx - cut_w // 2, 0, W) bby1 = np.clip(cy - cut_h // 2, 0, H) bbx2 = np.clip(cx + cut_w // 2, 0, W) bby2 = np.clip(cy + cut_h // 2, 0, H) return bbx1, bby1, bbx2, bby2

5.3 模型集成与知识蒸馏

提升模型性能的高级技巧：

class ModelEnsemble: """模型集成策略""" def __init__(self, models, weights=None): self.models = models if weights is None: self.weights = [1.0/len(models)] * len(models) else: self.weights = weights def predict(self, x): predictions = [] for model in self.models: model.eval() with torch.no_grad(): pred = model(x) predictions.append(pred) # 加权平均 ensemble_pred = sum(w * p for w, p in zip(self.weights, predictions)) return ensemble_pred class KnowledgeDistillation: """知识蒸馏训练""" def __init__(self, teacher_model, student_model, temperature=3, alpha=0.7): self.teacher = teacher_model self.student = student_model self.temperature = temperature self.alpha = alpha self.kl_loss = nn.KLDivLoss() def distill_loss(self, student_logits, teacher_logits, labels): # 教师模型软标签 soft_targets = F.softmax(teacher_logits / self.temperature, dim=-1) soft_prob = F.log_softmax(student_logits / self.temperature, dim=-1) # 蒸馏损失 distill_loss = self.kl_loss(soft_prob, soft_targets) * (self.temperature ** 2) # 学生模型硬标签损失 student_loss = F.cross_entropy(student_logits, labels) # 组合损失 return self.alpha * distill_loss + (1 - self.alpha) * student_loss

6. 完整项目实战与部署

6.1 端到端训练管道

整合所有组件的完整训练流程：

class RemoteSensingPipeline: """遥感深度学习端到端管道""" def __init__(self, config): self.config = config self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.setup_pipeline() def setup_pipeline(self): # 数据加载 self.train_loader, self.val_loader = self.create_dataloaders() # 模型初始化 self.model = self.create_model() # 优化器配置 self.optimizer = create_optimizer(self.model, self.config['optimizer']) self.criterion = self.create_criterion() # 训练记录 self.best_score = 0 self.train_history = { 'train_loss': [], 'val_loss': [], 'train_iou': [], 'val_iou': [] } def train(self): for epoch in range(self.config['epochs']): # 训练阶段 train_loss, train_iou = self.train_epoch() # 验证阶段 val_loss, val_iou = self.validate() # 学习率调整 self.scheduler.step(val_loss) # 记录历史 self.train_history['train_loss'].append(train_loss) self.train_history['val_loss'].append(val_loss) self.train_history['train_iou'].append(train_iou) self.train_history['val_iou'].append(val_iou) # 保存最佳模型 if val_iou > self.best_score: self.best_score = val_iou self.save_checkpoint(epoch, True) print(f'Epoch {epoch+1}: Train Loss: {train_loss:.4f}, ' f'Val Loss: {val_loss:.4f}, Val IoU: {val_iou:.4f}') def save_checkpoint(self, epoch, is_best=False): checkpoint = { 'epoch': epoch, 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'best_score': self.best_score, 'config': self.config } torch.save(checkpoint, f'checkpoint_epoch_{epoch}.pth') if is_best: torch.save(checkpoint, 'best_model.pth')

6.2 模型部署与推理优化

生产环境部署的优化策略：

class ModelDeployer: """模型部署工具类""" @staticmethod def convert_to_onnx(model, dummy_input, onnx