import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
import time
import os
def train_model(model, train_loader, valid_loader, criterion, optimizer, scheduler=None,
num_epochs=10, device=None, save_dir='./checkpoints'):
"""
训练模型并记录性能指标
参数:
model: 要训练的模型
train_loader, valid_loader: 训练和验证数据加载器
criterion: 损失函数
optimizer: 优化器
scheduler: 学习率调度器（可选）
num_epochs: 训练轮数
device: 使用的设备
save_dir: 模型保存目录
返回:
history: 包含训练历史的字典
"""
if device is None:
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
start_time = time.time()
model = model.to(device)
history = {
'train_loss': [], 'train_acc': [],
'val_loss': [], 'val_acc': [],
'epoch_times': []
}
best_val_acc = 0.0
# 确保保存目录存在
os.makedirs(save_dir, exist_ok=True)
for epoch in range(num_epochs):
epoch_start = time.time()
print(f"Epoch {epoch+1}/{num_epochs}")
# 训练阶段
model.train()
train_loss = 0.0
train_correct = 0
train_total = 0
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
# 梯度清零
optimizer.zero_grad()
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, labels)
# 反向传播和优化
loss.backward()
optimizer.step()
# 统计
train_loss += loss.item() * inputs.size(0)
_, predicted = torch.max(outputs, 1)
train_total += labels.size(0)
train_correct += (predicted == labels).sum().item()
# 计算训练指标
train_loss = train_loss / len(train_loader.sampler)
train_acc = train_correct / train_total
# 验证阶段
model.eval()
val_loss = 0.0
val_correct = 0
val_total = 0
with torch.no_grad():
for inputs, labels in valid_loader:
inputs, labels = inputs.to(device), labels.to(device)
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, labels)
# 统计
val_loss += loss.item() * inputs.size(0)
_, predicted = torch.max(outputs, 1)
val_total += labels.size(0)
val_correct += (predicted == labels).sum().item()
# 计算验证指标
val_loss = val_loss / len(valid_loader.sampler)
val_acc = val_correct / val_total
# 更新学习率
if scheduler:
scheduler.step()
# 记录历史
history['train_loss'].append(train_loss)
history['train_acc'].append(train_acc)
history['val_loss'].append(val_loss)
history['val_acc'].append(val_acc)
# 记录每个epoch的时间
epoch_end = time.time()
epoch_time = epoch_end - epoch_start
history['epoch_times'].append(epoch_time)
# 如果是最佳模型，保存权重
if val_acc > best_val_acc:
best_val_acc = val_acc
torch.save(model.state_dict(), f"{save_dir}/{model.__class__.__name__}_best.pth")
print(f"模型已保存到 {save_dir}/{model.__class__.__name__}_best.pth")
print(f"训练损失: {train_loss:.4f}, 训练准确率: {train_acc:.4f}")
print(f"验证损失: {val_loss:.4f}, 验证准确率: {val_acc:.4f}")
print(f"本轮用时: {epoch_time:.2f}s")
print("-" * 50)
# 计算总训练时间
total_time = time.time() - start_time
print(f"总训练时间: {total_time:.2f}s")
return model, history
def evaluate_model(model, test_loader, criterion, device=None, classes=None):
"""
评估模型在测试集上的性能
参数:
model: 要评估的模型
test_loader: 测试数据加载器
criterion: 损失函数
device: 使用的设备
classes: 类别名称列表
返回:
test_loss: 测试损失
test_acc: 测试准确率
"""
if device is None:
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
model.eval()
test_loss = 0.0
test_correct = 0
test_total = 0
y_true = []
y_pred = []
with torch.no_grad():
for inputs, labels in test_loader:
inputs, labels = inputs.to(device), labels.to(device)
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, labels)
# 统计
test_loss += loss.item() * inputs.size(0)
_, predicted = torch.max(outputs, 1)
test_total += labels.size(0)
test_correct += (predicted == labels).sum().item()
# 收集真实标签和预测标签
y_true.extend(labels.cpu().numpy())
y_pred.extend(predicted.cpu().numpy())
# 计算测试指标
test_loss = test_loss / len(test_loader.dataset)
test_acc = test_correct / test_total
print(f"测试损失: {test_loss:.4f}, 测试准确率: {test_acc:.4f}")
# 如果提供了类别名称，计算混淆矩阵
if classes:
try:
from sklearn.metrics import confusion_matrix, classification_report
import seaborn as sns
cm = confusion_matrix(y_true, y_pred)
plt.figure(figsize=(10, 8))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=classes, yticklabels=classes)
plt.xlabel('Predicted') # 英文标签
plt.ylabel('True') # 英文标签
plt.title('Confusion Matrix') # 英文标题
plt.show()
# 打印分类报告
print("分类报告:")
print(classification_report(y_true, y_pred, target_names=classes))
except ImportError:
print("警告: 未安装sklearn或seaborn，无法生成混淆矩阵和分类报告")
return test_loss, test_acc
def plot_training_history(history, title="Training History"):
"""
绘制训练历史曲线
参数:
history: 包含训练历史的字典
title: 图表标题
"""
plt.figure(figsize=(12, 5))
# 绘制损失曲线
plt.subplot(1, 2, 1)
plt.plot(history['train_loss'], label='Training Loss') # 英文标签
plt.plot(history['val_loss'], label='Validation Loss') # 英文标签
plt.xlabel('Epochs') # 英文标签
plt.ylabel('Loss') # 英文标签
plt.title('Loss Curves') # 英文标题
plt.legend()
# 绘制准确率曲线
plt.subplot(1, 2, 2)
plt.plot(history['train_acc'], label='Training Accuracy') # 英文标签
plt.plot(history['val_acc'], label='Validation Accuracy') # 英文标签
plt.xlabel('Epochs') # 英文标签
plt.ylabel('Accuracy') # 英文标签
plt.title('Accuracy Curves') # 英文标题
plt.legend()
plt.suptitle(title) # 英文总标题
plt.tight_layout()
plt.savefig(f"{title.replace(' ', '_')}.png")
plt.show()
def visualize_model_predictions(model, test_loader, classes, device=None, num_images=25):
"""
可视化模型预测
参数:
model: 要评估的模型
test_loader: 测试数据加载器
classes: 类别名称列表
device: 使用的设备
num_images: 要显示的图像数量
"""
if device is None:
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
model.eval()
# 获取batch数据
images, labels = next(iter(test_loader))
with torch.no_grad():
outputs = model(images.to(device))
_, preds = torch.max(outputs, 1)
# 将预测和标签转换为CPU上的numpy数组
preds = preds.cpu().numpy()
labels = labels.numpy()
# 计算display_grid的尺寸
grid_size = int(np.ceil(np.sqrt(num_images)))
fig, axes = plt.subplots(grid_size, grid_size, figsize=(15, 15))
for i, ax in enumerate(axes.flat):
if i < min(num_images, len(preds)):
img = images[i].numpy().transpose((1, 2, 0))
# 反标准化
mean = np.array([0.4914, 0.4822, 0.4465])
std = np.array([0.2023, 0.1994, 0.2010])
img = std * img + mean
img = np.clip(img, 0, 1)
ax.imshow(img)
color = "green" if preds[i] == labels[i] else "red"
ax.set_title(f"Predicted: {classes[preds[i]]}\nTrue: {classes[labels[i]]}", color=color) # 英文标签
ax.axis('off')
plt.tight_layout()
plt.show()
def visualize_conv_filters(model, layer_name='conv1'):
"""
可视化卷积核
参数:
model: 模型
layer_name: 要可视化的卷积层名称
"""
model.eval()
# 获取指定层的权重
for name, module in model.named_modules():
if name == layer_name and isinstance(module, nn.Conv2d):
weights = module.weight.data.clone().cpu()
break
else:
print(f"未找到名为 {layer_name} 的卷积层")
return
# 规范化权重以便可视化
weights = weights - weights.min()
weights = weights / weights.max()
# 绘制卷积核
num_filters = min(16, weights.size(0))
fig, axes = plt.subplots(4, 4, figsize=(10, 10))
fig.suptitle(f'Conv Layer {layer_name} Filters') # 英文标题
for i, ax in enumerate(axes.flat):
if i < num_filters:
# 如果是3通道的卷积核，直接显示RGB
if weights.size(1) == 3:
ax.imshow(weights[i].permute(1, 2, 0))
else:
# 如果不是3通道，只显示第一个通道
ax.imshow(weights[i, 0], cmap='viridis')
ax.axis('off')
plt.tight_layout()
plt.show()
def model_complexity(model, input_size=(3, 32, 32), batch_size=128, device=None):
"""
计算模型参数量和推理时间
参数:
model: 要评估的模型
input_size: 输入尺寸
batch_size: 批量大小
device: 使用的设备
返回:
num_params: 参数量
inference_time: 每批次推理时间
"""
if device is None:
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
model.eval()
# 计算参数量
num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
# 创建随机输入
dummy_input = torch.randn(batch_size, *input_size).to(device)
# 预热
with torch.no_grad():
for _ in range(10):
_ = model(dummy_input)
# 计时
start_time = time.time()
with torch.no_grad():
for _ in range(100):
_ = model(dummy_input)
end_time = time.time()
inference_time = (end_time - start_time) / 100
print(f"参数量: {num_params:,}")
print(f"每批次({batch_size}个样本)推理时间: {inference_time*1000:.2f}ms")
return num_params, inference_time