### PyTorch 实现简单猫狗图像分类模型 以下是基于 PyTorch 的猫狗图像分类模型的实现方法。该模型采用卷积神经网络 (Convolutional Neural Network, CNN)，并使用 Kaggle 提供的数据集进行训练。 #### 数据准备 首先，需要从 Kaggle 下载猫狗分类数据集[^3]。下载完成后，解压缩文件并将 `train` 文件夹中的图片按类别分开存储到两个子文件夹中（例如 `cats/` 和 `dogs/`）。这样可以方便后续加载器读取数据。 #### 安装依赖库 确保安装了必要的 Python 库： ```bash pip install torch torchvision matplotlib ``` #### 导入所需模块 在脚本开头导入所需的模块： ```python import os import torch from torch import nn, optim from torch.utils.data import DataLoader from torchvision import datasets, transforms from PIL import ImageFile ImageFile.LOAD_TRUNCATED_IMAGES = True ``` #### 数据预处理与增强 定义用于训练和验证的数据转换操作： ```python 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]) ]) } ``` 创建数据加载器： ```python data_dir = './data' # 替换为实际路径 image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val']} dataloaders = {x: DataLoader(image_datasets[x], batch_size=32, shuffle=True) for x in ['train', 'val']} dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} class_names = image_datasets['train'].classes ``` #### 构建 CNN 模型 设计一个简单的 CNN 结构来完成二分类任务： ```python class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.conv_layers = nn.Sequential( nn.Conv2d(3, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2) ) self.fc_layers = nn.Sequential( nn.Linear(64 * 56 * 56, 128), nn.ReLU(), nn.Dropout(p=0.5), nn.Linear(128, 2) # 输出层有两个节点分别对应猫和狗 ) def forward(self, x): x = self.conv_layers(x) x = x.view(-1, 64 * 56 * 56) x = self.fc_layers(x) return x ``` 初始化模型、损失函数以及优化器： ```python device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = SimpleCNN().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) ``` #### 训练过程 编写训练循环逻辑： ```python def train_model(model, criterion, optimizer, num_epochs=10): best_acc = 0.0 for epoch in range(num_epochs): print(f'Epoch {epoch}/{num_epochs - 1}') print('-' * 10) 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, labels = inputs.to(device), 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 / dataset_sizes[phase] epoch_acc = running_corrects.double() / dataset_sizes[phase] print(f'{phase} Loss: {epoch_loss:.4f}, Acc: {epoch_acc:.4f}') if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc print(f'Best val Acc: {best_acc:.4f}') return model ``` 调用训练函数启动训练流程： ```python trained_model = train_model(model, criterion, optimizer, num_epochs=25) torch.save(trained_model.state_dict(), 'cat_dog_classifier.pth') # 保存模型权重 ``` #### 预测阶段 预测新图像是猫还是狗时，需先加载已保存的模型参数，并设置评估模式： ```python def predict_image(img_path, model): transform = data_transforms['val'] img = Image.open(img_path).convert('RGB') img_t = transform(img) batch_t = torch.unsqueeze(img_t, 0) model.eval() out = model(batch_t.to(device)) _, index = torch.max(out, 1) percentage = torch.nn.functional.softmax(out, dim=1)[0] * 100 result = class_names[index] return f"{result}: {percentage[index].item():.2f}%" # 调用示例 print(predict_image('./test_images/dog.jpg', trained_model)) # 替换为测试图片的实际路径 ``` ---