pytorch cnn 实现猫狗分类

青丶空゛2025-02-18 15:43

### 文章目录

  • [@[toc]](#文章目录 @[toc] 1. 导入必要的库 2. 定义数据集类 3. 数据预处理和加载 4. 定义 CNN 模型 5. 定义损失函数和优化器 6. 训练模型 7. 保存模型 8. 使用模型进行预测 9 完整代码 10. 总结)
  • [1. 导入必要的库](#文章目录 @[toc] 1. 导入必要的库 2. 定义数据集类 3. 数据预处理和加载 4. 定义 CNN 模型 5. 定义损失函数和优化器 6. 训练模型 7. 保存模型 8. 使用模型进行预测 9 完整代码 10. 总结)
  • [2. 定义数据集类](#文章目录 @[toc] 1. 导入必要的库 2. 定义数据集类 3. 数据预处理和加载 4. 定义 CNN 模型 5. 定义损失函数和优化器 6. 训练模型 7. 保存模型 8. 使用模型进行预测 9 完整代码 10. 总结)
  • [3. 数据预处理和加载](#文章目录 @[toc] 1. 导入必要的库 2. 定义数据集类 3. 数据预处理和加载 4. 定义 CNN 模型 5. 定义损失函数和优化器 6. 训练模型 7. 保存模型 8. 使用模型进行预测 9 完整代码 10. 总结)
  • [4. 定义 CNN 模型](#文章目录 @[toc] 1. 导入必要的库 2. 定义数据集类 3. 数据预处理和加载 4. 定义 CNN 模型 5. 定义损失函数和优化器 6. 训练模型 7. 保存模型 8. 使用模型进行预测 9 完整代码 10. 总结)
  • [5. 定义损失函数和优化器](#文章目录 @[toc] 1. 导入必要的库 2. 定义数据集类 3. 数据预处理和加载 4. 定义 CNN 模型 5. 定义损失函数和优化器 6. 训练模型 7. 保存模型 8. 使用模型进行预测 9 完整代码 10. 总结)
  • [6. 训练模型](#文章目录 @[toc] 1. 导入必要的库 2. 定义数据集类 3. 数据预处理和加载 4. 定义 CNN 模型 5. 定义损失函数和优化器 6. 训练模型 7. 保存模型 8. 使用模型进行预测 9 完整代码 10. 总结)
  • [7. 保存模型](#文章目录 @[toc] 1. 导入必要的库 2. 定义数据集类 3. 数据预处理和加载 4. 定义 CNN 模型 5. 定义损失函数和优化器 6. 训练模型 7. 保存模型 8. 使用模型进行预测 9 完整代码 10. 总结)
  • [8. 使用模型进行预测](#文章目录 @[toc] 1. 导入必要的库 2. 定义数据集类 3. 数据预处理和加载 4. 定义 CNN 模型 5. 定义损失函数和优化器 6. 训练模型 7. 保存模型 8. 使用模型进行预测 9 完整代码 10. 总结)
  • [9 完整代码](#文章目录 @[toc] 1. 导入必要的库 2. 定义数据集类 3. 数据预处理和加载 4. 定义 CNN 模型 5. 定义损失函数和优化器 6. 训练模型 7. 保存模型 8. 使用模型进行预测 9 完整代码 10. 总结)
  • [10. 总结](#文章目录 @[toc] 1. 导入必要的库 2. 定义数据集类 3. 数据预处理和加载 4. 定义 CNN 模型 5. 定义损失函数和优化器 6. 训练模型 7. 保存模型 8. 使用模型进行预测 9 完整代码 10. 总结)

1. 导入必要的库

py 复制代码 
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import os

2. 定义数据集类

我们将创建一个自定义数据集类来加载猫狗图片。

py 复制代码 
class CatDogDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.classes = ['cat', 'dog']
        self.image_paths = []
        self.labels = []

        # 遍历 cat 和 dog 目录，加载图片路径和标签
        for idx, class_name in enumerate(self.classes):
            class_dir = os.path.join(root_dir, class_name)
            for img_name in os.listdir(class_dir):
                self.image_paths.append(os.path.join(class_dir, img_name))
                self.labels.append(idx)

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image_path = self.image_paths[idx]
        image = Image.open(image_path).convert('RGB')  # 确保图片是 RGB 格式
        label = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        return image, label

3. 数据预处理和加载

定义数据预处理方法，并加载数据集。

py 复制代码 
# 数据预处理
transform = transforms.Compose([
    transforms.Resize((64, 64)),  # 调整图片大小
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化
])

# 加载数据集
train_dataset = CatDogDataset(root_dir='path_to_train_data', transform=transform)
val_dataset = CatDogDataset(root_dir='path_to_val_data', transform=transform)
# 数据加载器

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

4. 定义 CNN 模型

py 复制代码 
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fc1 = nn.Linear(64 * 16 * 16, 512)
        self.fc2 = nn.Linear(512, 2)
        self.dropout = nn.Dropout(0.5)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 64 * 16 * 16)
        x = self.dropout(self.relu(self.fc1(x)))
        x = self.fc2(x)
        return x

model = CNN()

5. 定义损失函数和优化器

py 复制代码 
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

6. 训练模型

py 复制代码 
num_epochs = 10
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
    
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}")

    # 验证模型
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    print(f"Validation Accuracy: {100 * correct / total:.2f}%")

7. 保存模型

py 复制代码 
torch.save(model.state_dict(), 'cat_dog_classifier.pth')

8. 使用模型进行预测

py 复制代码 
# 加载模型
model.load_state_dict(torch.load('cat_dog_classifier.pth'))
model.eval()

# 预测函数
def predict_image(image_path):
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0).to(device)
    
    with torch.no_grad():
        output = model(image)
        _, predicted = torch.max(output, 1)
    
    return 'cat' if predicted.item() == 0 else 'dog'

# 使用模型进行预测
image_path = 'path_to_test_image.jpg'
prediction = predict_image(image_path)
print(f"The image is a {prediction}")

# 使用模型进行预测
image_path = 'path_to_test_image.jpg'
prediction = predict_image(image_path)
print(f"The image is a {prediction}")

9 完整代码

py 复制代码 
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import os

class CatDogDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.classes = ['cats', 'dogs']
        self.image_paths = []
        self.labels = []

        # 遍历 cat 和 dog 目录，加载图片路径和标签
        for idx, class_name in enumerate(self.classes):
            class_dir = os.path.join(root_dir, class_name)
            num_pets = 0
            for img_name in os.listdir(class_dir):
                self.image_paths.append(os.path.join(class_dir, img_name))
                self.labels.append(idx)
                # print("class_dir : ", img_name)
                # num_pets = num_pets + 1
                # if num_pets >= 5000:
                #     break


    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image_path = self.image_paths[idx]
        image = Image.open(image_path).convert('RGB')  # 确保图片是 RGB 格式
        label = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        return image, label

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((64, 64)),  # 调整图片大小
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化
])

# 加载数据集
train_dataset = CatDogDataset(root_dir='D:/Cache/dataset/PetImages/train', transform=transform)
val_dataset = CatDogDataset(root_dir='D:/Cache/dataset/PetImages/valid', transform=transform)

# 数据加载器
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fc1 = nn.Linear(64 * 16 * 16, 512)
        self.fc2 = nn.Linear(512, 2)
        self.dropout = nn.Dropout(0.5)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 64 * 16 * 16)
        x = self.dropout(self.relu(self.fc1(x)))
        x = self.fc2(x)
        return x

model = CNN()

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def train():
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
        
        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}")

        # 验证模型
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        
        print(f"Validation Accuracy: {100 * correct / total:.2f}%")


    torch.save(model.state_dict(), 'cat_dog_classifier.pth')

# 预测函数
def predict_image(image_path):
    
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0).to(device)
    
    with torch.no_grad():
        output = model(image)
        _, predicted = torch.max(output, 1)
    
    return 'cat' if predicted.item() == 0 else 'dog'


def test():
    # 使用模型进行预测
    # 加载模型
    model.load_state_dict(torch.load('cat_dog_classifier.pth'))
    model.eval()
    image_path = 'D:/Cache/dataset/PetImages/Dog/6.jpg'
    image_path = 'D:/develop/pytorch/dogcat/img/training/dogs/dog1.jpg'
    image_path = 'D:/develop/pytorch/dogcat/img/training/cats/4.jpg'
    prediction = predict_image(image_path)
    print(f"The image is a {prediction}")

import matplotlib.pyplot as plt

def test1():
    image_path = 'D:/develop/pytorch/dogcat/img/training/cats/3.jpg'
    img = Image.open(image_path).convert('RGB')

    transform = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
    ])

    img = transform(img)
    img = img.unsqueeze(0)  # 添加batch维度

    model.load_state_dict(torch.load('cat_dog_classifier.pth'))
    model.eval()
    prediction = predict_image(image_path)

    class_names = ['cat', 'dog']
    print("Predicted class:", prediction)

    plt.imshow(img.squeeze().numpy().transpose((1, 2, 0)))
    plt.show()

if __name__ == '__main__':
    train()
    test1()
复制代码 
D:\develop\pytorch\dogcat>python3.7 dogVsCat.py
C:\Users\yosola\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.7_qbz5n2kfra8p0\LocalCache\local-packages\Python37\site-packages\PIL\TiffImagePlugin.py:864: UserWarning: Truncated File Read
  warnings.warn(str(msg))
Epoch [1/10], Loss: 0.5782
Validation Accuracy: 76.25%
Epoch [2/10], Loss: 0.4676
Validation Accuracy: 81.10%
Epoch [3/10], Loss: 0.4201
Validation Accuracy: 84.90%
Epoch [4/10], Loss: 0.3605
Validation Accuracy: 88.25%
Epoch [5/10], Loss: 0.2949
Validation Accuracy: 92.50%
Epoch [6/10], Loss: 0.2234
Validation Accuracy: 95.90%
Epoch [7/10], Loss: 0.1562
Validation Accuracy: 98.00%
Epoch [8/10], Loss: 0.1069
Validation Accuracy: 98.60%
Epoch [9/10], Loss: 0.0907
Validation Accuracy: 99.70%
Epoch [10/10], Loss: 0.0785
Validation Accuracy: 99.50%
Predicted class: cat

10. 总结

我们定义了一个自定义数据集类 CatDogDataset 来加载猫狗图片。

使用 PyTorch 的 DataLoader 加载数据。

定义了一个简单的 CNN 模型进行训练。

保存训练好的模型，并使用模型进行预测。

你可以根据需要调整模型的架构、超参数和数据增强方法。希望这个示例对你有帮助！

相关推荐
YYXZZ。。1 小时前
PyTorch——搭建小实战和Sequential的使用（7）
人工智能·pytorch·python
四川兔兔1 小时前
pytorch 与 张量的处理
人工智能·pytorch·python
yvestine6 小时前
自然语言处理——文本分类
自然语言处理·分类·文本分类·评价指标·pr·roc
SpikeKing7 小时前
Server - 使用 Docker 配置 PyTorch 研发环境
pytorch·docker·llm
大鹏的NLP博客13 小时前
基于 Transformer robert的情感分类任务实践总结之二——R-Drop
分类·transformer·r-dop
Humbunklung14 小时前
机器学习算法分类
算法·机器学习·分类
weixin_4487816214 小时前
DenseNet算法 实现乳腺癌识别
pytorch·深度学习·神经网络
2301_8050545615 小时前
Python训练营打卡Day48(2025.6.8）
pytorch·python·深度学习
王上上17 小时前
【论文阅读28】-CNN-BiLSTM-Attention-（2024）
论文阅读·人工智能·cnn
只有左边一个小酒窝19 小时前
（六）卷积神经网络：深度学习在计算机视觉中的应用
深度学习·计算机视觉·cnn
热门推荐
01【图像处理与机器视觉】XJTU期末考点02从零安装 LLaMA-Factory 微调 Qwen 大模型成功及所有的坑03KGG转MP3工具|非KGM文件|解密音频04海康Visionmaster-常见问题排查方法-启动阶段05YOLOv8入门 | 重要性能衡量指标、训练结果评价及分析及影响mAP的因素【发论文关注的指标】06【SpeedAI科研小助手】2分钟极速解决知网维普重复率、AIGC率过高，一键全文降！文件格式不变，公式都保留的！07Coze扣子平台完整体验和实践（附国内和国际版对比）08DeepSeek各版本说明与优缺点分析09VMware虚拟机安装Win7专业版保姆级教程（附镜像包）10R-tree详解