根据你的需求，我会完善代码以适配你提供的图片路径，并确保代码能够正常运行。以下是完整代码和详细解释。 --- ### 完善后的代码 ```python import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms from torch.utils.data import DataLoader, Dataset import os from PIL import Image import matplotlib.pyplot as plt import numpy as np from torch.utils.tensorboard import SummaryWriter from datetime import datetime # 1. 数据集准备和预处理 class CustomDataset(Dataset): def __init__(self, root_dir, transform=None): """ 自定义数据集类 :param root_dir: 数据集根目录 :param transform: 数据预处理变换 """ self.root_dir = root_dir self.transform = transform self.classes = sorted(os.listdir(root_dir)) self.class_to_idx = {cls_name: i for i, cls_name in enumerate(self.classes)} self.images = [] # 遍历每个类别文件夹，收集图像路径和标签 for cls_name in self.classes: cls_dir = os.path.join(root_dir, cls_name) if not os.path.isdir(cls_dir): # 确保是文件夹 continue for img_name in os.listdir(cls_dir): if img_name.endswith(('.jpg', '.jpeg', '.png')): self.images.append((os.path.join(cls_dir, img_name), self.class_to_idx[cls_name])) def __len__(self): return len(self.images) def __getitem__(self, idx): img_path, label = self.images[idx] image = Image.open(img_path).convert('RGB') if self.transform: image = self.transform(image) return image, label # 2. 定义AlexNet8模型 class AlexNet8(nn.Module): def __init__(self, num_classes=10): """ AlexNet8模型定义 :param num_classes: 分类类别数 """ super(AlexNet8, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), nn.Conv2d(64, 192, kernel_size=5, padding=2), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), nn.Conv2d(192, 384, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(384, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(256, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), ) self.avgpool = nn.AdaptiveAvgPool2d((6, 6)) self.classifier = nn.Sequential( nn.Dropout(), nn.Linear(256 * 6 * 6, 4096), nn.ReLU(inplace=True), nn.Dropout(), nn.Linear(4096, 4096), nn.ReLU(inplace=True), nn.Linear(4096, num_classes), ) def forward(self, x): x = self.features(x) x = self.avgpool(x) x = torch.flatten(x, 1) x = self.classifier(x) return x # 3. 训练和评估函数 def train_model(model, criterion, optimizer, scheduler, dataloaders, dataset_sizes, device, num_epochs=25, log_dir='runs'): """ 模型训练函数 :param model: 要训练的模型 :param criterion: 损失函数 :param optimizer: 优化器 :param scheduler: 学习率调度器 :param dataloaders: 包含训练和验证数据的数据加载器字典 :param dataset_sizes: 数据集大小字典 :param device: 训练设备 (CPU/GPU) :param num_epochs: 训练轮数 :param log_dir: TensorBoard日志目录 :return: 训练好的模型 """ since = datetime.now() writer = SummaryWriter(log_dir=os.path.join(log_dir, f'alexnet8_{since.strftime("%Y%m%d_%H%M%S")}')) best_model_wts = model.state_dict() 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) if phase == 'train': scheduler.step() epoch_loss = running_loss / dataset_sizes[phase] epoch_acc = running_corrects.double() / dataset_sizes[phase] # 记录到TensorBoard writer.add_scalar(f'Loss/{phase}', epoch_loss, epoch) writer.add_scalar(f'Accuracy/{phase}', epoch_acc, epoch) print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}') # 深度复制模型 if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = model.state_dict() print() time_elapsed = datetime.now() - since print(f'Training complete in {time_elapsed}') print(f'Best val Acc: {best_acc:4f}') # 加载最佳模型权重 model.load_state_dict(best_model_wts) writer.close() return model def visualize_predictions(model, dataloader, class_names, device, num_images=9): """ 可视化模型预测结果 :param model: 训练好的模型 :param dataloader: 数据加载器 :param class_names: 类别名称列表 :param device: 设备 (CPU/GPU) :param num_images: 要显示的图像数量 """ model.eval() images_so_far = 0 plt.figure(figsize=(10, 10)) with torch.no_grad(): for inputs, labels in dataloader: inputs = inputs.to(device) labels = labels.to(device) outputs = model(inputs) _, preds = torch.max(outputs, 1) for j in range(inputs.size()[0]): images_so_far += 1 ax = plt.subplot(3, 3, images_so_far) ax.axis('off') ax.set_title(f'True: {class_names[labels[j]]}\nPred: {class_names[preds[j]]}') # 反标准化图像 img = inputs.cpu().data[j].numpy().transpose((1, 2, 0)) mean = np.array([0.485, 0.456, 0.406]) std = np.array([0.229, 0.224, 0.225]) img = std * img + mean img = np.clip(img, 0, 1) ax.imshow(img) if images_so_far == num_images: plt.tight_layout() plt.show() return # 4. 主函数 def main(): # 设置随机种子以保证可重复性 torch.manual_seed(42) np.random.seed(42) # 检查GPU是否可用 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # 数据预处理和增强 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 = r'C:\Users\周程锦\PycharmProjects\pythonProject5\dl\ceshi' image_datasets = { 'train': CustomDataset(os.path.join(data_dir, 'train'), data_transforms['train']), 'val': CustomDataset(os.path.join(data_dir, 'val'), data_transforms['val']) } # 创建数据加载器 batch_size = 32 dataloaders = { 'train': DataLoader(image_datasets['train'], batch_size=batch_size, shuffle=True, num_workers=0), 'val': DataLoader(image_datasets['val'], batch_size=batch_size, shuffle=False, num_workers=0) } dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} class_names = image_datasets['train'].classes print(f"Dataset sizes: {dataset_sizes}") print(f"Class names: {class_names}") # 初始化模型 model = AlexNet8(num_classes=len(class_names)) model = model.to(device) # 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9) # 每7个epoch衰减一次学习率 exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1) # 训练模型 model = train_model(model, criterion, optimizer, exp_lr_scheduler, dataloaders, dataset_sizes, device, num_epochs=25) # 保存模型 torch.save(model.state_dict(), 'alexnet8_model.pth') print("Model saved to alexnet8_model.pth") # 加载模型 loaded_model = AlexNet8(num_classes=len(class_names)) loaded_model.load_state_dict(torch.load('alexnet8_model.pth')) loaded_model = loaded_model.to(device) print("Model loaded from alexnet8_model.pth") # 可视化预测结果 visualize_predictions(loaded_model, dataloaders['val'], class_names, device) if __name__ == '__main__': main() ``` --- ### 代码解释 #### 1. 数据集路径修改 - 将 `data_dir` 设置为你提供的路径： ```python data_dir = r'C:\Users\周程锦\PycharmProjects\pythonProject5\dl\ceshi' ``` - 确保该路径下有 `train` 和 `val` 文件夹，分别存放训练集和验证集。 #### 2. 图像加载 - 使用 `PIL.Image.open` 打开图像，并通过 `transforms` 对图像进行预处理（如缩放、裁剪、归一化）。 #### 3. 模型训练 - 在 `train_model` 函数中，模型会自动训练并记录 TensorBoard 日志。 - 每个 epoch 结束后，保存验证集上表现最好的模型权重。 #### 4. 结果可视化 - 使用 `visualize_predictions` 函数展示模型的预测结果，并与真实标签对比。 --- ### 相关问题