写一个完整的模型代码，我有五类图像数据，已经放入文件夹并且命名为test train val，请训练我的数据，并且将权重保存到某个文件

<think>我们根据用户的需求，构建一个基于CBAM-ResNet的模型，用于训练五类图像数据，并保存训练后的模型权重。 我们将结合之前的讨论，特别是CBAM模块的实现和ResNet的结构。同时，注意用户提供的引用中关于CBAM模块的一些细节（如kernel_size为奇数，以保证空间维度不变）[^1]。 我们将构建一个CBAM-ResNet模型，这里以ResNet18为基础，在残差块中加入CBAM模块。由于是五分类任务，我们将最后的全连接层输出改为5。 步骤： 1. 定义CBAM模块（通道注意力和空间注意力）。 2. 定义CBAMResNetBlock，将CBAM模块插入到残差块中。 3. 构建CBAM-ResNet网络，这里我们构建一个类似ResNet18的结构，由4个阶段组成，每个阶段包含2个残差块（第一个残差块可能进行下采样）。 4. 定义训练过程，包括数据加载、模型训练、验证和保存模型权重。 注意：由于用户引用[1]中提到kernel_size为7（奇数）可以保持空间尺寸不变，我们在空间注意力模块中使用kernel_size=7。 我们假设输入图像大小为224x224（这是ResNet常用的输入尺寸），但实际中可以根据数据调整。 代码结构： - 导入必要的库（torch, torchvision, torch.nn等） - 定义ChannelAttention模块 - 定义SpatialAttention模块 - 定义CBAMResNetBlock（基于ResNet的基本残差块，加入CBAM） - 定义CBAMResNet（整个网络） - 训练函数和主程序 由于用户要求用于五类图像数据，我们假设数据已经准备好，放在train和val文件夹中，每个类一个子文件夹。 由于代码较长，我们将分步实现。同时，为了简化，我们只实现ResNet18的结构（即每个阶段2个残差块，深度为[2,2,2,2]）。 注意：在残差块中，CBAM模块放在两个卷积层之后，残差连接之前（即对残差块的主路径输出进行注意力加权，然后再与shortcut相加）。 另外，我们使用预训练模型可以加速收敛，但这里我们从头训练（如果用户有预训练需求，可以后续添加，但用户要求五分类，所以需要修改最后的全连接层，因此我们这里不使用预训练）。 开始编写代码：</think>以下是一个完整的基于 PyTorch 实现的 CBAM-ResNet 模型代码，用于训练五类图像分类任务并保存模型权重： ```python import torch import torch.nn as nn import torch.optim as optim import torchvision.transforms as transforms from torch.utils.data import DataLoader, Dataset from torchvision.datasets import ImageFolder import os from PIL import Image # CBAM 模块实现 class CBAMBlock(nn.Module): def __init__(self, channel, reduction=16, kernel_size=7): super(CBAMBlock, self).__init__() # 通道注意力 self.avg_pool = nn.AdaptiveAvgPool2d(1) self.max_pool = nn.AdaptiveMaxPool2d(1) self.fc = nn.Sequential( nn.Linear(channel, channel // reduction, bias=False), nn.ReLU(inplace=True), nn.Linear(channel // reduction, channel, bias=False) ) self.sigmoid = nn.Sigmoid() # 空间注意力 - 使用奇数kernel_size确保空间尺寸不变[^1] self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False) def forward(self, x): b, c, _, _ = x.size() # 通道注意力 avg_out = self.fc(self.avg_pool(x).view(b, c)) max_out = self.fc(self.max_pool(x).view(b, c)) channel_att = self.sigmoid(avg_out + max_out).view(b, c, 1, 1) x = x * channel_att # 空间注意力 avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) spatial_att = torch.cat([avg_out, max_out], dim=1) spatial_att = self.sigmoid(self.conv(spatial_att)) return x * spatial_att # 残差块集成CBAM class CBAMResBlock(nn.Module): def __init__(self, in_channels, out_channels, stride=1): super(CBAMResBlock, 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.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(out_channels) self.relu = nn.ReLU(inplace=True) # CBAM注意力模块 self.cbam = CBAMBlock(channel=out_channels, reduction=16, kernel_size=7) # 下采样 self.downsample = None if stride != 1 or in_channels != out_channels: self.downsample = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(out_channels) ) def forward(self, x): identity = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) # 应用CBAM注意力 out = self.cbam(out) if self.downsample is not None: identity = self.downsample(x) out += identity out = self.relu(out) return out # CBAM-ResNet模型 class CBAMResNet(nn.Module): def __init__(self, block, layers, num_classes=5): super(CBAMResNet, self).__init__() self.in_channels = 64 # 初始卷积层 self.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(block, 64, layers[0], stride=1) self.layer2 = self._make_layer(block, 128, layers[1], stride=2) self.layer3 = self._make_layer(block, 256, layers[2], stride=2) self.layer4 = self._make_layer(block, 512, layers[3], stride=2) # 分类器 self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(512, num_classes) def _make_layer(self, block, out_channels, blocks, stride=1): layers = [] layers.append(block(self.in_channels, out_channels, stride)) self.in_channels = out_channels for _ in range(1, blocks): layers.append(block(out_channels, out_channels)) return nn.Sequential(*layers) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.maxpool(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.avgpool(x) x = torch.flatten(x, 1) x = self.fc(x) return x # 创建CBAM-ResNet18模型 def cbam_resnet18(num_classes=5): return CBAMResNet(CBAMResBlock, [2, 2, 2, 2], num_classes=num_classes) # 训练函数 def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, model_save_path='cbam_resnet.pth'): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = model.to(device) best_acc = 0.0 for epoch in range(num_epochs): print(f'Epoch {epoch}/{num_epochs-1}') print('-' * 10) # 每个epoch都有训练和验证阶段 for phase in ['train', 'val']: if phase == 'train': model.train() else: model.eval() running_loss = 0.0 running_corrects = 0 # 迭代数据 for inputs, labels in dataloaders[phase]: inputs = inputs.to(device) labels = labels.to(device) # 梯度清零 optimizer.zero_grad() # 前向传播 with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) # 训练阶段反向传播+优化 if phase == 'train': loss.backward() optimizer.step() # 统计 running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / len(dataloaders[phase].dataset) epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset) print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}') # 保存最佳模型 if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc torch.save(model.state_dict(), model_save_path) print(f'Saved best model with val acc: {best_acc:.4f}') return model # 主程序 if __name__ == '__main__': # 数据预处理 data_transforms = { 'train': transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } # 数据集路径（根据实际情况修改） data_dir = 'path/to/your/dataset' # 包含train和val子文件夹 image_datasets = { 'train': ImageFolder(os.path.join(data_dir, 'train'), data_transforms['train']), 'val': ImageFolder(os.path.join(data_dir, 'val'), data_transforms['val']) } # 创建数据加载器 dataloaders = { 'train': DataLoader(image_datasets['train'], batch_size=32, shuffle=True, num_workers=4), 'val': DataLoader(image_datasets['val'], batch_size=32, shuffle=False, num_workers=4) } # 初始化模型 model = cbam_resnet18(num_classes=5) # 损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) # 训练模型 trained_model = train_model( model, dataloaders, criterion, optimizer, num_epochs=50, model_save_path='cbam_resnet_best.pth' ) # 保存最终模型 torch.save(trained_model.state_dict(), 'cbam_resnet_final.pth') print("Training complete and models saved!") ``` ### 使用说明： 1. **数据集准备**： - 创建文件夹结构： ``` your_dataset/ ├── train/ │ ├── class1/ │ ├── class2/ │ ├── class3/ │ ├── class4/ │ └── class5/ └── val/ ├── class1/ ├── class2/ ├── class3/ ├── class4/ └── class5/ ``` - 将代码中的 `data_dir = 'path/to/your/dataset'` 修改为实际路径 2. **模型特点**： - 基于ResNet18架构，在每个残差块后添加CBAM注意力模块 - 通道注意力使用降维因子16（reduction=16） - 空间注意力使用7×7卷积核（kernel_size=7）确保空间尺寸不变[^1] - 输出层适配五分类任务 3. **训练配置**： - 使用Adam优化器（学习率0.001） - 交叉熵损失函数 - 数据增强：随机裁剪、水平翻转 - 输入图像尺寸：224×224 - 批量大小：32 4. **模型保存**： - 最佳模型：`cbam_resnet_best.pth`（基于验证集准确率） - 最终模型：`cbam_resnet_final.pth` 5. **扩展应用**： - 可修改为ResNet50/101等更深结构 - 可迁移到目标检测任务（如YOLOv4）[^2] - 可与其他骨干网络结合（如MobileNet）[^2] ### 训练建议： 1. 使用GPU加速训练（如可用） 2. 监控训练/验证损失曲线防止过拟合 3. 学习率可考虑使用调度器（如StepLR） 4. 对于小数据集，可使用迁移学习（加载预训练权重）[^4] ###