Python训练营打卡Day49

作业：尝试对今天的模型检查参数数目，并用tensorboard查看训练过程

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
import torchvision.models as models
from torchsummary import summary
import numpy as np
import os# 通道注意力模块
class ChannelAttention(nn.Module):def __init__(self, in_channels, reduction_ratio=16):super(ChannelAttention, self).__init__()self.avg_pool = nn.AdaptiveAvgPool2d(1)self.max_pool = nn.AdaptiveMaxPool2d(1)self.fc = nn.Sequential(nn.Conv2d(in_channels, in_channels // reduction_ratio, 1, bias=False),nn.ReLU(),nn.Conv2d(in_channels // reduction_ratio, in_channels, 1, bias=False))self.sigmoid = nn.Sigmoid()def forward(self, x):avg_out = self.fc(self.avg_pool(x))max_out = self.fc(self.max_pool(x))out = avg_out + max_outreturn self.sigmoid(out)# 空间注意力模块
class SpatialAttention(nn.Module):def __init__(self, kernel_size=7):super(SpatialAttention, self).__init__()self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False)self.sigmoid = nn.Sigmoid()def forward(self, x):avg_out = torch.mean(x, dim=1, keepdim=True)max_out, _ = torch.max(x, dim=1, keepdim=True)x_cat = torch.cat([avg_out, max_out], dim=1)out = self.conv(x_cat)return self.sigmoid(out)# CBAM模块
class CBAM(nn.Module):def __init__(self, in_channels, reduction_ratio=16, kernel_size=7):super(CBAM, self).__init__()self.channel_attention = ChannelAttention(in_channels, reduction_ratio)self.spatial_attention = SpatialAttention(kernel_size)def forward(self, x):x = x * self.channel_attention(x)x = x * self.spatial_attention(x)return x# 带CBAM的ResNet模块
class ResidualBlockWithCBAM(nn.Module):def __init__(self, in_channels, out_channels, stride=1, reduction_ratio=16):super(ResidualBlockWithCBAM, self).__init__()self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)self.bn1 = nn.BatchNorm2d(out_channels)self.relu = nn.ReLU(inplace=True)self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)self.bn2 = nn.BatchNorm2d(out_channels)self.cbam = CBAM(out_channels, reduction_ratio)self.shortcut = nn.Sequential()if stride != 1 or in_channels != out_channels:self.shortcut = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),nn.BatchNorm2d(out_channels))def forward(self, x):residual = xout = self.relu(self.bn1(self.conv1(x)))out = self.bn2(self.conv2(out))# 应用CBAMout = self.cbam(out)out += self.shortcut(residual)out = self.relu(out)return out# 简化版ResNet18+CBAM模型
class ResNet18WithCBAM(nn.Module):def __init__(self, num_classes=1000, reduction_ratio=16):super(ResNet18WithCBAM, self).__init__()self.in_channels = 64self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)self.bn1 = nn.BatchNorm2d(64)self.relu = nn.ReLU(inplace=True)self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)# 残差块self.layer1 = self._make_layer(64, 2, stride=1, reduction_ratio=reduction_ratio)self.layer2 = self._make_layer(128, 2, stride=2, reduction_ratio=reduction_ratio)self.layer3 = self._make_layer(256, 2, stride=2, reduction_ratio=reduction_ratio)self.layer4 = self._make_layer(512, 2, stride=2, reduction_ratio=reduction_ratio)self.avgpool = nn.AdaptiveAvgPool2d((1, 1))self.fc = nn.Linear(512, num_classes)def _make_layer(self, out_channels, num_blocks, stride, reduction_ratio):strides = [stride] + [1] * (num_blocks - 1)layers = []for stride in strides:layers.append(ResidualBlockWithCBAM(self.in_channels, out_channels, stride, reduction_ratio))self.in_channels = out_channelsreturn nn.Sequential(*layers)def forward(self, x):out = self.relu(self.bn1(self.conv1(x)))out = self.maxpool(out)out = self.layer1(out)out = self.layer2(out)out = self.layer3(out)out = self.layer4(out)out = self.avgpool(out)out = torch.flatten(out, 1)out = self.fc(out)return out# 计算模型参数数量
def count_parameters(model):return sum(p.numel() for p in model.parameters() if p.requires_grad)# 训练函数
def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=10, device='cuda'):# 创建TensorBoard记录器writer = SummaryWriter('runs/cbam_experiment')model.to(device)# 添加模型图到TensorBoardfor images, _ in train_loader:images = images.to(device)writer.add_graph(model, images)breakbest_val_acc = 0.0for epoch in range(num_epochs):# 训练阶段model.train()train_loss = 0.0train_correct = 0train_total = 0for batch_idx, (inputs, targets) in enumerate(train_loader):inputs, targets = inputs.to(device), targets.to(device)optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, targets)loss.backward()optimizer.step()train_loss += loss.item()_, predicted = outputs.max(1)train_total += targets.size(0)train_correct += predicted.eq(targets).sum().item()# 每100个batch记录一次损失if (batch_idx + 1) % 100 == 0:writer.add_scalar('Training Loss', train_loss / 100, epoch * len(train_loader) + batch_idx)train_loss = 0.0# 验证阶段model.eval()val_loss = 0.0val_correct = 0val_total = 0with torch.no_grad():for inputs, targets in val_loader:inputs, targets = inputs.to(device), targets.to(device)outputs = model(inputs)loss = criterion(outputs, targets)val_loss += loss.item()_, predicted = outputs.max(1)val_total += targets.size(0)val_correct += predicted.eq(targets).sum().item()# 计算准确率train_acc = 100.0 * train_correct / train_totalval_acc = 100.0 * val_correct / val_total# 记录准确率到TensorBoardwriter.add_scalar('Training Accuracy', train_acc, epoch)writer.add_scalar('Validation Accuracy', val_acc, epoch)# 记录学习率writer.add_scalar('Learning Rate', optimizer.param_groups[0]['lr'], epoch)# 保存最佳模型if val_acc > best_val_acc:best_val_acc = val_acctorch.save(model.state_dict(), 'best_cbam_model.pth')print(f'Epoch: {epoch+1}/{num_epochs}, Train Acc: {train_acc:.2f}%, Val Acc: {val_acc:.2f}%')writer.close()return model# 示例：如何使用上述代码
def main():# 创建模型model_with_cbam = ResNet18WithCBAM(num_classes=10, reduction_ratio=16)model_vanilla = models.resnet18(pretrained=False)model_vanilla.fc = nn.Linear(512, 10)  # 修改最后一层以匹配类别数# 计算参数数量params_with_cbam = count_parameters(model_with_cbam)params_vanilla = count_parameters(model_vanilla)print(f"ResNet18 with CBAM parameters: {params_with_cbam:,}")print(f"Vanilla ResNet18 parameters: {params_vanilla:,}")print(f"Parameter increase: {((params_with_cbam - params_vanilla) / params_vanilla) * 100:.2f}%")# 这里应该有数据加载代码，为了简化，我们只展示训练函数的调用# train_loader, val_loader = load_data()# 定义损失函数和优化器# criterion = nn.CrossEntropyLoss()# optimizer = torch.optim.SGD(model_with_cbam.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)# 训练模型# trained_model = train_model(model_with_cbam, train_loader, val_loader, criterion, optimizer, num_epochs=5)if __name__ == "__main__":main()

@浙大疏锦行