311 lines
11 KiB
Python
311 lines
11 KiB
Python
import os
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.optim as optim
|
||
from torch.utils.data import DataLoader
|
||
import numpy as np
|
||
from tqdm import tqdm
|
||
import matplotlib.pyplot as plt
|
||
|
||
from .models import get_model
|
||
from .data import RemoteSensingDataset, get_train_transforms, get_val_transforms
|
||
from .utils.metrics import calculate_metrics
|
||
from .utils.losses import get_loss_function
|
||
|
||
class UAVSegTrainer:
|
||
|
||
def __init__(self,
|
||
data_dir,
|
||
model_type='deeplabv3plus',
|
||
num_classes=6,
|
||
batch_size=4,
|
||
epochs=100,
|
||
learning_rate=0.001,
|
||
tile_size=512,
|
||
stride=None,
|
||
loss_type='ce',
|
||
device=None,
|
||
save_dir='./checkpoints'):
|
||
"""
|
||
Args:
|
||
data_dir (str): 数据目录，应包含'images'和'masks'子目录
|
||
model_type (str): 模型类型，可选 'deeplabv3plus' 或 'unetpp'
|
||
num_classes (int): 类别数量
|
||
batch_size (int): 批次大小
|
||
epochs (int): 训练轮数
|
||
learning_rate (float): 学习率
|
||
tile_size (int): 分块大小
|
||
stride (int): 滑动窗口步长，如果为None则等于tile_size
|
||
loss_type (str): 损失函数类型，可选 'ce'(交叉熵), 'dice', 'focal', 'combined'
|
||
device (str): 设备，如果为None则自动选择
|
||
save_dir (str): 模型保存目录
|
||
"""
|
||
self.data_dir = data_dir
|
||
self.model_type = model_type
|
||
self.num_classes = num_classes
|
||
self.batch_size = batch_size
|
||
self.epochs = epochs
|
||
self.learning_rate = learning_rate
|
||
self.tile_size = tile_size
|
||
self.stride = stride
|
||
self.loss_type = loss_type
|
||
self.save_dir = save_dir
|
||
|
||
# 设置设备
|
||
if device is None:
|
||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
else:
|
||
self.device = torch.device(device)
|
||
|
||
print(f"使用设备: {self.device}")
|
||
|
||
# 创建保存目录
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
|
||
# 初始化模型
|
||
self.model = get_model(model_type, num_classes)
|
||
self.model.to(self.device)
|
||
|
||
# 初始化损失函数
|
||
self.criterion = get_loss_function(loss_type, num_classes)
|
||
|
||
# 初始化优化器
|
||
self.optimizer = optim.Adam([
|
||
{'params': self.model.get_backbone_params(), 'lr': learning_rate * 0.1},
|
||
{'params': self.model.get_decoder_params(), 'lr': learning_rate}
|
||
])
|
||
|
||
# 学习率调度器
|
||
self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
|
||
self.optimizer, mode='min', factor=0.5, patience=5, verbose=True
|
||
)
|
||
|
||
# 训练历史记录
|
||
self.history = {
|
||
'train_loss': [],
|
||
'val_loss': [],
|
||
'metrics': {}
|
||
}
|
||
|
||
def _get_dataloaders(self):
|
||
"""准备数据加载器"""
|
||
# 数据目录
|
||
train_img_dir = os.path.join(self.data_dir, 'train', 'images')
|
||
train_mask_dir = os.path.join(self.data_dir, 'train', 'masks')
|
||
val_img_dir = os.path.join(self.data_dir, 'val', 'images')
|
||
val_mask_dir = os.path.join(self.data_dir, 'val', 'masks')
|
||
|
||
# 数据增强和预处理
|
||
train_transform = get_train_transforms()
|
||
val_transform = get_val_transforms()
|
||
|
||
# 训练数据集
|
||
train_dataset = RemoteSensingDataset(
|
||
image_dir=train_img_dir,
|
||
mask_dir=train_mask_dir,
|
||
transform=train_transform,
|
||
tile_size=self.tile_size,
|
||
stride=self.stride
|
||
)
|
||
|
||
# 验证数据集
|
||
val_dataset = RemoteSensingDataset(
|
||
image_dir=val_img_dir,
|
||
mask_dir=val_mask_dir,
|
||
transform=val_transform,
|
||
tile_size=self.tile_size,
|
||
stride=self.tile_size # 验证时不重叠
|
||
)
|
||
|
||
# 数据加载器
|
||
train_loader = DataLoader(
|
||
train_dataset,
|
||
batch_size=self.batch_size,
|
||
shuffle=True,
|
||
num_workers=0,
|
||
pin_memory=True
|
||
)
|
||
|
||
val_loader = DataLoader(
|
||
val_dataset,
|
||
batch_size=self.batch_size,
|
||
shuffle=False,
|
||
num_workers=0,
|
||
pin_memory=True
|
||
)
|
||
|
||
return train_loader, val_loader
|
||
|
||
def _train_epoch(self, train_loader):
|
||
"""训练一个轮次"""
|
||
self.model.train()
|
||
epoch_loss = 0
|
||
|
||
with tqdm(train_loader, desc="训练") as pbar:
|
||
for images, masks in pbar:
|
||
images = images.to(self.device)
|
||
masks = masks.to(self.device)
|
||
|
||
# 前向传播
|
||
outputs = self.model(images)
|
||
loss = self.criterion(outputs, masks)
|
||
|
||
# 反向传播和优化
|
||
self.optimizer.zero_grad()
|
||
loss.backward()
|
||
self.optimizer.step()
|
||
|
||
# 更新进度条
|
||
epoch_loss += loss.item()
|
||
pbar.set_postfix({'loss': loss.item()})
|
||
|
||
return epoch_loss / len(train_loader)
|
||
|
||
def _validate(self, val_loader):
|
||
"""验证模型"""
|
||
self.model.eval()
|
||
val_loss = 0
|
||
metrics_sum = None
|
||
samples_count = 0
|
||
|
||
with torch.no_grad():
|
||
with tqdm(val_loader, desc="验证") as pbar:
|
||
for images, masks in pbar:
|
||
images = images.to(self.device)
|
||
masks = masks.to(self.device)
|
||
batch_size = images.size(0)
|
||
|
||
# 前向传播
|
||
outputs = self.model(images)
|
||
loss = self.criterion(outputs, masks)
|
||
|
||
# 计算指标
|
||
batch_metrics = calculate_metrics(outputs, masks, self.num_classes)
|
||
|
||
# 累加损失和指标
|
||
val_loss += loss.item() * batch_size
|
||
if metrics_sum is None:
|
||
metrics_sum = {k: v * batch_size for k, v in batch_metrics.items()}
|
||
else:
|
||
for k, v in batch_metrics.items():
|
||
metrics_sum[k] += v * batch_size
|
||
|
||
samples_count += batch_size
|
||
|
||
# 更新进度条
|
||
pbar.set_postfix({'loss': loss.item()})
|
||
|
||
# 计算平均值
|
||
val_loss /= samples_count
|
||
metrics = {k: v / samples_count for k, v in metrics_sum.items()}
|
||
|
||
return val_loss, metrics
|
||
|
||
def train(self):
|
||
"""训练模型"""
|
||
# 准备数据加载器
|
||
train_loader, val_loader = self._get_dataloaders()
|
||
|
||
# 最佳验证损失
|
||
best_val_loss = float('inf')
|
||
|
||
# 训练循环
|
||
for epoch in range(1, self.epochs + 1):
|
||
print(f"\n轮次 {epoch}/{self.epochs}")
|
||
|
||
# 训练
|
||
train_loss = self._train_epoch(train_loader)
|
||
self.history['train_loss'].append(train_loss)
|
||
|
||
# 验证
|
||
val_loss, metrics = self._validate(val_loader)
|
||
self.history['val_loss'].append(val_loss)
|
||
|
||
# 更新学习率
|
||
self.scheduler.step(val_loss)
|
||
|
||
# 保存指标
|
||
for k, v in metrics.items():
|
||
if k not in self.history['metrics']:
|
||
self.history['metrics'][k] = []
|
||
self.history['metrics'][k].append(v)
|
||
|
||
# 打印指标
|
||
print(f"训练损失: {train_loss:.4f}, 验证损失: {val_loss:.4f}")
|
||
for k, v in metrics.items():
|
||
print(f"{k}: {v:.4f}")
|
||
|
||
# 保存最佳模型
|
||
if val_loss < best_val_loss:
|
||
best_val_loss = val_loss
|
||
self._save_model(f"{self.model_type}_best.pth")
|
||
print(f"保存最佳模型，验证损失: {val_loss:.4f}")
|
||
|
||
# 保存最新模型
|
||
if epoch % 10 == 0:
|
||
self._save_model(f"{self.model_type}_epoch_{epoch}.pth")
|
||
|
||
# 保存最终模型
|
||
self._save_model(f"{self.model_type}_final.pth")
|
||
|
||
# 绘制训练曲线
|
||
self._plot_training_curves()
|
||
|
||
return self.history
|
||
|
||
def _save_model(self, filename):
|
||
"""保存模型"""
|
||
save_path = os.path.join(self.save_dir, filename)
|
||
torch.save({
|
||
'model_state_dict': self.model.state_dict(),
|
||
'optimizer_state_dict': self.optimizer.state_dict(),
|
||
'model_type': self.model_type,
|
||
'num_classes': self.num_classes,
|
||
}, save_path)
|
||
|
||
def _plot_training_curves(self):
|
||
"""绘制训练曲线"""
|
||
# 创建图形
|
||
plt.figure(figsize=(15, 5))
|
||
#设置宋体
|
||
plt.rcParams['font.sans-serif'] = ['SimSun']
|
||
plt.rcParams['axes.unicode_minus'] = False
|
||
# 绘制损失曲线
|
||
plt.subplot(1, 2, 1)
|
||
plt.plot(self.history['train_loss'], label='训练损失')
|
||
plt.plot(self.history['val_loss'], label='验证损失')
|
||
plt.xlabel('轮次')
|
||
plt.ylabel('损失')
|
||
plt.legend()
|
||
plt.title('训练和验证损失')
|
||
|
||
# 绘制指标曲线
|
||
plt.subplot(1, 2, 2)
|
||
for metric_name, metric_values in self.history['metrics'].items():
|
||
plt.plot(metric_values, label=metric_name)
|
||
plt.xlabel('轮次')
|
||
plt.ylabel('指标值')
|
||
plt.legend()
|
||
plt.title('验证指标')
|
||
|
||
# 保存图形
|
||
plt.tight_layout()
|
||
plt.savefig(os.path.join(self.save_dir, f"{self.model_type}_training_curves.png"))
|
||
plt.close()
|
||
|
||
def load_model(self, checkpoint_path):
|
||
"""加载模型"""
|
||
checkpoint = torch.load(checkpoint_path, map_location=self.device)
|
||
|
||
# 检查模型类型
|
||
if checkpoint['model_type'] != self.model_type:
|
||
print(f"警告: 加载的模型类型 ({checkpoint['model_type']}) 与当前模型类型 ({self.model_type}) 不匹配")
|
||
|
||
# 加载模型参数
|
||
self.model.load_state_dict(checkpoint['model_state_dict'])
|
||
|
||
# 加载优化器参数（如果存在）
|
||
if 'optimizer_state_dict' in checkpoint:
|
||
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
|
||
|
||
print(f"模型已从 {checkpoint_path} 加载") |