python
复制代码
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np
from torch.optim.lr_scheduler import ReduceLROnPlateau, CosineAnnealingLR
# 设置中文字体支持
plt.rcParams["font.family"] = ["SimHei"]
plt.rcParams['axes.unicode_minus'] = False # 解决负号显示问题
# 1. 增强的数据预处理和数据增强
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4), # 随机裁剪
transforms.RandomHorizontalFlip(p=0.5), # 随机水平翻转
transforms.RandomRotation(15), # 随机旋转
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), # 颜色抖动
transforms.ToTensor(), # 转换为张量
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)) # CIFAR-10标准归一化
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])
# 2. 加载CIFAR-10数据集
train_dataset = datasets.CIFAR10(
root='./data',
train=True,
download=True,
transform=transform_train
)
test_dataset = datasets.CIFAR10(
root='./data',
train=False,
transform=transform_test
)
# 3. 创建数据加载器（调整batch_size）
batch_size = 128 # 增大batch_size
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
# 4. 改进的MLP模型（添加批归一化和更多层）
class ImprovedMLP(nn.Module):
def __init__(self, hidden_layers=[1024, 512, 256], dropout_rate=0.3):
super(ImprovedMLP, self).__init__()
self.flatten = nn.Flatten()
layers = []
input_size = 3072 # 3*32*32
# 动态构建隐藏层
for i, hidden_size in enumerate(hidden_layers):
layers.append(nn.Linear(input_size, hidden_size))
layers.append(nn.BatchNorm1d(hidden_size)) # 添加批归一化
layers.append(nn.ReLU(inplace=True))
layers.append(nn.Dropout(dropout_rate))
input_size = hidden_size
# 输出层
layers.append(nn.Linear(input_size, 10))
self.network = nn.Sequential(*layers)
def forward(self, x):
return self.network(self.flatten(x))
# 检查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")
# 初始化模型（选择其中一个）
model = ImprovedMLP(hidden_layers=[2048, 1024, 512, 256], dropout_rate=0.4)
# model = MLPWithResidual()
model = model.to(device)
# 参数统计
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"总参数数量: {total_params:,}")
print(f"可训练参数数量: {trainable_params:,}")
# 使用标签平滑的损失函数
class LabelSmoothingCrossEntropy(nn.Module):
def __init__(self, smoothing=0.1):
super(LabelSmoothingCrossEntropy, self).__init__()
self.smoothing = smoothing
def forward(self, x, target):
log_probs = torch.nn.functional.log_softmax(x, dim=-1)
nll_loss = -log_probs.gather(dim=-1, index=target.unsqueeze(1))
nll_loss = nll_loss.squeeze(1)
smooth_loss = -log_probs.mean(dim=-1)
loss = (1 - self.smoothing) * nll_loss + self.smoothing * smooth_loss
return loss.mean()
# 选择损失函数
criterion = LabelSmoothingCrossEntropy(smoothing=0.1) # 或使用 nn.CrossEntropyLoss()
# 优化器调参
optimizer = optim.AdamW(model.parameters(),
lr=0.001,
weight_decay=1e-4, # 权重衰减
betas=(0.9, 0.999))
# 学习率调度器
scheduler = CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6)
# scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=3, verbose=True)
# 6. 增强的训练函数
def train(model, train_loader, test_loader, criterion, optimizer, scheduler, device, epochs):
model.train()
train_losses = []
test_losses = []
train_accuracies = []
test_accuracies = []
all_iter_losses = []
iter_indices = []
best_accuracy = 0.0
for epoch in range(epochs):
# 训练阶段
model.train()
running_loss = 0.0
correct = 0
total = 0
# 使用混合精度训练（如果可用）
scaler = torch.cuda.amp.GradScaler() if device.type == 'cuda' else None
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
if scaler:
# 混合精度训练
with torch.cuda.amp.autocast():
output = model(data)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
# 记录损失
iter_loss = loss.item()
all_iter_losses.append(iter_loss)
iter_indices.append(epoch * len(train_loader) + batch_idx + 1)
running_loss += iter_loss
_, predicted = output.max(1)
total += target.size(0)
correct += predicted.eq(target).sum().item()
# 梯度裁剪
if scaler:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
else:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
if (batch_idx + 1) % 50 == 0:
acc = 100. * correct / total
print(f'Epoch: {epoch+1}/{epochs} | Batch: {batch_idx+1}/{len(train_loader)} '
f'| 损失: {iter_loss:.4f} | 准确率: {acc:.2f}%')
# 计算训练统计
epoch_train_loss = running_loss / len(train_loader)
epoch_train_acc = 100. * correct / total
train_losses.append(epoch_train_loss)
train_accuracies.append(epoch_train_acc)
# 测试阶段
model.eval()
test_loss = 0
correct_test = 0
total_test = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += criterion(output, target).item()
_, predicted = output.max(1)
total_test += target.size(0)
correct_test += predicted.eq(target).sum().item()
epoch_test_loss = test_loss / len(test_loader)
epoch_test_acc = 100. * correct_test / total_test
test_losses.append(epoch_test_loss)
test_accuracies.append(epoch_test_acc)
# 更新学习率
if isinstance(scheduler, ReduceLROnPlateau):
scheduler.step(epoch_test_acc)
else:
scheduler.step()
# 保存最佳模型
if epoch_test_acc > best_accuracy:
best_accuracy = epoch_test_acc
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'accuracy': best_accuracy,
}, 'best_model.pth')
# 打印epoch总结
current_lr = optimizer.param_groups[0]['lr']
print(f'Epoch {epoch+1}/{epochs} 完成 | '
f'训练损失: {epoch_train_loss:.4f} | 训练准确率: {epoch_train_acc:.2f}% | '
f'测试损失: {epoch_test_loss:.4f} | 测试准确率: {epoch_test_acc:.2f}% | '
f'学习率: {current_lr:.6f}')
print('-' * 80)
# 绘制结果
plot_results(train_losses, test_losses, train_accuracies, test_accuracies, all_iter_losses, iter_indices)
return best_accuracy
# 7. 绘制综合结果
def plot_results(train_losses, test_losses, train_accs, test_accs, iter_losses, iter_indices):
fig, axes = plt.subplots(2, 2, figsize=(15, 10))
# 1. 迭代损失
axes[0, 0].plot(iter_indices, iter_losses, 'b-', alpha=0.5, linewidth=0.5)
axes[0, 0].set_xlabel('Iteration（Batch序号）')
axes[0, 0].set_ylabel('损失值')
axes[0, 0].set_title('每个 Iteration 的训练损失')
axes[0, 0].grid(True, alpha=0.3)
# 2. 每个epoch的训练和测试损失
epochs = range(1, len(train_losses) + 1)
axes[0, 1].plot(epochs, train_losses, 'b-', label='训练损失')
axes[0, 1].plot(epochs, test_losses, 'r-', label='测试损失')
axes[0, 1].set_xlabel('Epoch')
axes[0, 1].set_ylabel('损失值')
axes[0, 1].set_title('训练和测试损失')
axes[0, 1].legend()
axes[0, 1].grid(True, alpha=0.3)
# 3. 每个epoch的训练和测试准确率
axes[1, 0].plot(epochs, train_accs, 'b-', label='训练准确率')
axes[1, 0].plot(epochs, test_accs, 'r-', label='测试准确率')
axes[1, 0].set_xlabel('Epoch')
axes[1, 0].set_ylabel('准确率 (%)')
axes[1, 0].set_title('训练和测试准确率')
axes[1, 0].legend()
axes[1, 0].grid(True, alpha=0.3)
# 4. 损失和准确率的滑动平均
window_size = 100
smoothed_losses = np.convolve(iter_losses, np.ones(window_size)/window_size, mode='valid')
smoothed_indices = iter_indices[window_size-1:]
axes[1, 1].plot(smoothed_indices, smoothed_losses, 'g-', linewidth=2)
axes[1, 1].set_xlabel('Iteration（Batch序号）')
axes[1, 1].set_ylabel('损失值（滑动平均）')
axes[1, 1].set_title(f'训练损失（{window_size}次迭代滑动平均）')
axes[1, 1].grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
def train_single_fold(model, train_loader, val_loader, criterion, optimizer, scheduler, device, epochs, fold_num):
best_acc = 0.0
for epoch in range(epochs):
model.train()
for data, target in train_loader:
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
scheduler.step()
# 验证
model.eval()
correct = 0
total = 0
with torch.no_grad():
for data, target in val_loader:
data, target = data.to(device), target.to(device)
output = model(data)
_, predicted = output.max(1)
total += target.size(0)
correct += predicted.eq(target).sum().item()
val_acc = 100. * correct / total
if val_acc > best_acc:
best_acc = val_acc
print(f'Fold {fold_num} | Epoch {epoch+1}/{epochs} | 验证准确率: {val_acc:.2f}%')
return best_acc
# 9. 主执行部分
if __name__ == "__main__":
# 选择训练模式
training_mode = "normal" # 可选: "normal", "kfold"
if training_mode == "normal":
# 正常训练
epochs = 50
print("开始训练模型...")
best_accuracy = train(model, train_loader, test_loader, criterion, optimizer, scheduler, device, epochs)
print(f"训练完成！最佳测试准确率: {best_accuracy:.2f}%")
# 加载最佳模型并测试
checkpoint = torch.load('best_model.pth')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
# 最终测试
correct = 0
total = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
_, predicted = output.max(1)
total += target.size(0)
correct += predicted.eq(target).sum().item()
final_acc = 100. * correct / total
print(f"最终测试准确率: {final_acc:.2f}%")
