当前位置：首页 > news >正文

DAY 41 超大力王爱学Python

news 来源：原创 2025/6/4 11:41:14

DAY 41 简单CNN

数据增强
卷积神经网络定义的写法
batch归一化：调整一个批次的分布，常用与图像数据
特征图：只有卷积操作输出的才叫特征图
调度器：直接修改基础学习率

卷积操作常见流程如下：

1. 输入 → 卷积层 → Batch归一化层（可选） → 池化层 → 激活函数 → 下一层

Flatten -> Dense (with Dropout，可选) -> Dense (Output)

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
from typing import Tuple, Callable, Optional
import numpy as np# 1. 数据预处理与加载函数（增强版）
def get_data_loaders(dataset_name: str = 'MNIST',  # 可选: 'MNIST'或'CIFAR10'batch_size: int = 64,data_dir: str = './data',num_workers: int = 2,use_augmentation: bool = True  # 是否使用数据增强
) -> Tuple[DataLoader, DataLoader]:"""获取训练和测试数据加载器，支持灰度(MNIST)和彩色(CIFAR10)数据集"""# 根据数据集类型设置不同的转换if dataset_name == 'MNIST':# 基础转换base_transform = transforms.Compose([transforms.ToTensor(),  # 转换为Tensor并归一化到[0,1]transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差])# 增强转换augment_transform = transforms.Compose([transforms.RandomRotation(10),  # 随机旋转±10度transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),  # 随机平移transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,))])train_transform = augment_transform if use_augmentation else base_transformtrain_dataset = datasets.MNIST(data_dir, train=True, download=True, transform=train_transform)test_dataset = datasets.MNIST(data_dir, train=False, transform=base_transform)elif dataset_name == 'CIFAR10':# 基础转换base_transform = transforms.Compose([transforms.ToTensor(),  # 转换为Tensor并归一化到[0,1]transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 标准化到[-1,1]])# 增强转换augment_transform = transforms.Compose([transforms.RandomCrop(32, padding=4),  # 随机裁剪transforms.RandomHorizontalFlip(),  # 随机水平翻转transforms.RandomRotation(15),  # 随机旋转±15度transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])train_transform = augment_transform if use_augmentation else base_transformtrain_dataset = datasets.CIFAR10(data_dir, train=True, download=True, transform=train_transform)test_dataset = datasets.CIFAR10(data_dir, train=False, transform=base_transform)else:raise ValueError(f"Unsupported dataset: {dataset_name}")# 创建数据加载器train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,num_workers=num_workers)test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=False,num_workers=num_workers)return train_loader, test_loader# 2. 卷积神经网络定义（替代原有MLP）
class Flatten(nn.Module):"""自定义展平层，保留batch维度"""def forward(self, x):return x.view(x.size(0), -1)class CNNClassifier(nn.Module):"""卷积神经网络分类器，支持灰度和彩色图像"""def __init__(self,input_channels: int = 1,  # MNIST:1, CIFAR10:3num_classes: int = 10,use_batchnorm: bool = True  # 是否使用Batch归一化):super().__init__()# 定义特征提取部分（卷积层）if input_channels == 1:  # 灰度图像(MNIST)self.features = nn.Sequential(# 第一层卷积nn.Conv2d(input_channels, 32, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(32) if use_batchnorm else nn.Identity(),nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2),  # 输出: 14x14# 第二层卷积nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(64) if use_batchnorm else nn.Identity(),nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2),  # 输出: 7x7# 第三层卷积nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(128) if use_batchnorm else nn.Identity(),nn.ReLU()  # 输出: 7x7)# 展平后的特征维度flattened_size = 128 * 7 * 7else:  # 彩色图像(CIFAR10)self.features = nn.Sequential(# 第一层卷积nn.Conv2d(input_channels, 32, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(32) if use_batchnorm else nn.Identity(),nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2),  # 输出: 16x16# 第二层卷积nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(64) if use_batchnorm else nn.Identity(),nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2),  # 输出: 8x8# 第三层卷积nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(128) if use_batchnorm else nn.Identity(),nn.ReLU(),nn.MaxPool2d(kernel_size=2, stride=2)  # 输出: 4x4)# 展平后的特征维度flattened_size = 128 * 4 * 4# 定义分类部分（全连接层）self.classifier = nn.Sequential(Flatten(),nn.Dropout(0.5),nn.Linear(flattened_size, 128),nn.ReLU(),nn.Dropout(0.5),nn.Linear(128, num_classes))# 用于特征图可视化的钩子self.feature_maps = []self.register_forward_hook(self._save_feature_maps)def forward(self, x):x = self.features(x)x = self.classifier(x)return xdef _save_feature_maps(self, module, input, output):"""保存中间特征图，用于可视化"""if isinstance(module, nn.Conv2d):self.feature_maps.append(output.detach())# 3. 特征图可视化函数
def visualize_feature_maps(model, sample_image, device):"""可视化卷积层输出的特征图"""model.eval()sample_image = sample_image.to(device).unsqueeze(0)  # 添加batch维度# 前向传播并捕获特征图model.feature_maps = []  # 清空之前的特征图with torch.no_grad():model(sample_image)# 绘制原始图像plt.figure(figsize=(15, 10))plt.subplot(2, 3, 1)img = sample_image.cpu().squeeze(0).permute(1, 2, 0)img = img * torch.tensor([0.5, 0.5, 0.5]) + torch.tensor([0.5, 0.5, 0.5])  # 反归一化img = np.clip(img.numpy(), 0, 1)plt.imshow(img)plt.title('Original Image')plt.axis('off')# 绘制前三个卷积层的特征图for i, feature_map in enumerate(model.feature_maps[:3]):if i >= 5:  # 最多显示5个特征图break# 选择第一个样本的第一个通道plt.subplot(2, 3, i+2)plt.imshow(feature_map[0, 0].cpu(), cmap='viridis')plt.title(f'Feature Map {i+1}')plt.axis('off')plt.tight_layout()plt.savefig('feature_maps.png')plt.close()print("特征图已保存为 'feature_maps.png'")# 4. 训练函数（修改以支持学习率调度器）
def train(model: nn.Module,train_loader: DataLoader,criterion: nn.Module,optimizer: optim.Optimizer,scheduler: Optional[StepLR],  # 学习率调度器device: torch.device,epoch: int,log_interval: int = 100
) -> None:"""训练模型一个epoch"""model.train()  # 启用训练模式（激活dropout等）running_loss = 0.0for batch_idx, (data, target) in enumerate(train_loader):data, target = data.to(device), target.to(device)# 前向传播optimizer.zero_grad()output = model(data)loss = criterion(output, target)# 反向传播loss.backward()optimizer.step()running_loss += loss.item()# 打印训练进度if batch_idx % log_interval == 0:current_lr = optimizer.param_groups[0]['lr']print(f'Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} 'f'({100. * batch_idx / len(train_loader):.0f}%)]\t'f'Loss: {loss.item():.6f}\tLR: {current_lr:.6f}')# 打印平均损失avg_loss = running_loss / len(train_loader)print(f'Epoch {epoch} average loss: {avg_loss:.4f}')# 更新学习率if scheduler:scheduler.step()print(f'Epoch {epoch} learning rate updated to: {optimizer.param_groups[0]["lr"]:.6f}')# 5. 测试函数（保持不变）
def test(model: nn.Module,test_loader: DataLoader,criterion: nn.Module,device: torch.device
) -> Tuple[float, float]:"""评估模型在测试集上的性能"""model.eval()  # 启用评估模式（关闭dropout等）test_loss = 0correct = 0with 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()  # 累加批次损失pred = output.argmax(dim=1, keepdim=True)  # 获取最大概率的类别correct += pred.eq(target.view_as(pred)).sum().item()  # 统计正确预测数# 计算平均损失和准确率test_loss /= len(test_loader)accuracy = 100. * correct / len(test_loader.dataset)print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} 'f'({accuracy:.2f}%)\n')return test_loss, accuracy# 6. 主函数：训练和测试流程（增强版）
def main(dataset_name: str = 'MNIST',batch_size: int = 64,epochs: int = 5,lr: float = 0.001,gamma: float = 0.7,  # 学习率衰减因子step_size: int = 1,  # 学习率衰减步数use_augmentation: bool = True,use_batchnorm: bool = True,use_cuda: bool = True
) -> None:"""主函数：整合数据加载、模型训练和测试流程"""# 设置设备device = torch.device("cuda" if (use_cuda and torch.cuda.is_available()) else "cpu")print(f"Using device: {device}")# 获取数据加载器train_loader, test_loader = get_data_loaders(dataset_name=dataset_name,batch_size=batch_size,use_augmentation=use_augmentation)# 确定输入参数if dataset_name == 'MNIST':input_channels = 1num_classes = 10elif dataset_name == 'CIFAR10':input_channels = 3num_classes = 10else:raise ValueError(f"Unsupported dataset: {dataset_name}")# 初始化模型model = CNNClassifier(input_channels=input_channels,num_classes=num_classes,use_batchnorm=use_batchnorm).to(device)# 定义损失函数和优化器criterion = nn.CrossEntropyLoss()optimizer = optim.Adam(model.parameters(), lr=lr)# 定义学习率调度器scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma)# 训练和测试循环best_accuracy = 0.0for epoch in range(1, epochs + 1):train(model, train_loader, criterion, optimizer, scheduler, device, epoch)test_loss, accuracy = test(model, test_loader, criterion, device)# 保存最佳模型if accuracy > best_accuracy:best_accuracy = accuracytorch.save(model.state_dict(), f"{dataset_name}_best_cnn_model.pth")print(f"Best model saved with accuracy: {accuracy:.2f}%")# 可视化特征图（使用测试集中的第一个样本）sample_image, _ = next(iter(test_loader))visualize_feature_maps(model, sample_image[0], device)print(f"Training completed. Best accuracy: {best_accuracy:.2f}%")if __name__ == "__main__":# 训练MNIST模型main(dataset_name='MNIST',batch_size=64,epochs=5,lr=0.001,gamma=0.7,step_size=1,use_augmentation=True,use_batchnorm=True)

@浙大疏锦行