Pytorch 猫狗识别案例

猫狗识别数据集https://download.csdn.net/download/Victor_Li_/88483483?spm=1001.2014.3001.5501

训练集图片路径

测试集图片路径

训练代码如下

python 复制代码
import torch
import torchvision
import matplotlib.pyplot as plt
import torchvision.models as models
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp
import time
from torch.optim.lr_scheduler import StepLR

if __name__ == '__main__':
    torch.autograd.set_detect_anomaly(True)
    mp.freeze_support()
    train_on_gpu = torch.cuda.is_available()
    if not train_on_gpu:
        print('CUDA is not available. Training on CPU...')
    else:
        print('CUDA is available! Training on GPU...')

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    batch_size = 32
    # 设置数据预处理的转换
    transform = torchvision.transforms.Compose([
        torchvision.transforms.Resize((224, 224)),  # 调整图像大小为 224x224
        torchvision.transforms.RandomHorizontalFlip(),
        torchvision.transforms.RandomRotation(45),
        torchvision.transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
        torchvision.transforms.ToTensor(),  # 转换为张量
        torchvision.transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 归一化
    ])
    dataset = torchvision.datasets.ImageFolder('./cats_and_dogs_train',
                                               transform=transform)

    val_ratio = 0.2
    val_size = int(len(dataset) * val_ratio)
    train_size = len(dataset) - val_size
    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])

    train_dataset = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4,
                                                pin_memory=True)
    val_dataset = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, num_workers=4, pin_memory=True)

    # x,y = next(iter(val_dataset))
    # x = x.permute(1, 2, 0)  # 将通道维度调整到最后
    # x = (x - x.min()) / (x.max() - x.min())  # 反归一化操作
    # plt.imshow(x)  # 将通道维度调整到最后
    # plt.axis('off')  # 关闭坐标轴
    # plt.show()

    model = models.resnet34(weights=None)

    num_classes = 2
    model.fc = nn.Sequential(
        nn.Dropout(p=0.2),
        # nn.BatchNorm4d(model.fc.in_features),
        nn.Linear(model.fc.in_features, num_classes),
        nn.Sigmoid(),
    )
    lambda_L1 = 0.001
    lambda_L2 = 0.0001
    regularization_loss_L1 = 0
    regularization_loss_L2 = 0
    for name,param in model.named_parameters():
        param.requires_grad = True
        if 'bias' not in name:
            regularization_loss_L1 += torch.norm(param, p=1).detach()
            regularization_loss_L2 += torch.norm(param, p=2).detach()

    optimizer = optim.Adam(model.parameters(), lr=0.01)
    scheduler = StepLR(optimizer, step_size=5, gamma=0.9)
    criterion = nn.BCELoss().to(device)

    model.to(device)
    # print(model)
    loadfilename = "recognize_cats_and_dogs.pt"
    savefilename = "recognize_cats_and_dogs3.pt"

    checkpoint = torch.load(loadfilename)
    model.load_state_dict(checkpoint['model_state_dict'])


    def save_checkpoint(epoch, model, optimizer, filename, train_loss=0., val_loss=0.):
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'train_loss': train_loss,
            'val_loss': val_loss,
        }
        torch.save(checkpoint, filename)


    num_epochs = 100
    train_loss = []
    for epoch in range(num_epochs):
        running_loss = 0
        correct = 0
        total = 0
        epoch_start_time = time.time()
        for i, (inputs, labels) in enumerate(train_dataset):
            # 将数据放到设备上
            inputs, labels = inputs.to(device), labels.to(device)
            # 前向计算
            outputs = model(inputs)
            one_hot = nn.functional.one_hot(labels, num_classes).float()
            # 计算损失和梯度
            loss = criterion(outputs, one_hot) + lambda_L1 * regularization_loss_L1 + lambda_L2 * regularization_loss_L2
            loss.backward()
            if ((i + 1) % 2 == 0) or (i + 1 == len(train_dataset)):
                # 更新模型参数
                optimizer.step()
                optimizer.zero_grad()

            # 记录损失和准确率
            running_loss += loss.item()
            train_loss.append(loss.item())
            _, predicted = torch.max(outputs.data, 1)
            correct += (predicted == labels).sum().item()
            total += labels.size(0)
        accuracy_train = 100 * correct / total
        # 在测试集上计算准确率
        with torch.no_grad():
            running_loss_test = 0
            correct_test = 0
            total_test = 0
            for inputs, labels in val_dataset:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                one_hot = nn.functional.one_hot(labels, num_classes).float()
                loss = criterion(outputs, one_hot)
                running_loss_test += loss.item()

                _, predicted = torch.max(outputs.data, 1)
                correct_test += (predicted == labels).sum().item()
                total_test += labels.size(0)
            accuracy_test = 100 * correct_test / total_test
            # 输出每个 epoch 的损失和准确率
        epoch_end_time = time.time()
        epoch_time = epoch_end_time - epoch_start_time
        tain_loss = running_loss / len(train_dataset)
        val_loss = running_loss_test / len(val_dataset)
        print(
            "Epoch [{}/{}], Time: {:.4f}s, Loss: {:.4f}, Train Accuracy: {:.2f}%, Loss: {:.4f}, Test Accuracy: {:.2f}%"
            .format(epoch + 1, num_epochs, epoch_time, tain_loss,
                    accuracy_train, val_loss, accuracy_test))
        save_checkpoint(epoch, model, optimizer, savefilename, tain_loss, val_loss)
        scheduler.step()

    # plt.plot(train_loss, label='Train Loss')
    # # 添加图例和标签
    # plt.legend()
    # plt.xlabel('Epochs')
    # plt.ylabel('Loss')
    # plt.title('Training Loss')
    #
    # # 显示图形
    # plt.show()

测试代码如下

python 复制代码
import torch
import torchvision
import torch.nn as nn
import torchvision.models as models
import matplotlib.pyplot as plt
import torch.multiprocessing as mp

if __name__ == '__main__':
    mp.freeze_support()
    train_on_gpu = torch.cuda.is_available()
    if not train_on_gpu:
        print('CUDA is not available. Training on CPU...')
    else:
        print('CUDA is available! Training on GPU...')

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    batch_size = 32
    transform = torchvision.transforms.Compose([
        torchvision.transforms.Resize((224,224)),  # 调整图像大小为 224x224
        torchvision.transforms.ToTensor(),  # 转换为张量
        torchvision.transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 归一化
    ])
    dataset = torchvision.datasets.ImageFolder('./cats_and_dogs_test',
                                                     transform=transform)

    test_dataset = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True,num_workers=4, pin_memory=True)

    model = models.resnet34()

    num_classes = 2
    for param in model.parameters():
        param.requires_grad = False

    model.fc = nn.Sequential(
        nn.Dropout(),
        nn.Linear(model.fc.in_features,num_classes),
        nn.LogSoftmax(dim=1)
    )
    model.to(device)
    # print(model)

    filename = "recognize_cats_and_dogs.pt"
    checkpoint = torch.load(filename)
    model.load_state_dict(checkpoint['model_state_dict'])

    class_name = ['cat','dog']
    # 在测试集上计算准确率
    with torch.no_grad():
        for inputs, labels in test_dataset:
            inputs, labels = inputs.to(device), labels.to(device)
            output = model(inputs)
            _, predicted = torch.max(output.data, 1)
            for x,y,z in zip(inputs,labels,predicted):
                x = (x - x.min()) / (x.max() - x.min())
                plt.imshow(x.cpu().permute(1,2,0))
                plt.axis('off')
                plt.title('predicted: {0}'.format(class_name[z]))
                plt.show()

部分测试结果如下

相关推荐
酷雷曼VR全景5 分钟前
系统升级丨让VR全景制作更全面、更简单
人工智能·科技·vr·vr全景·酷雷曼
Goboy8 分钟前
Python修仙入门 - 踏入仙门的第一步
后端·python
链上日记15 分钟前
数字资产反诈指南:识别套路,守护WEEX账户安全
人工智能·安全
爱吃青菜的大力水手21 分钟前
软银收购 ABB 机器人业务的背景与 Physical AI 的关联
人工智能·机器人
长桥夜波42 分钟前
【第十七周】机器学习笔记06
人工智能·笔记·机器学习
程序员爱钓鱼44 分钟前
Python编程实战 · 基础入门篇 | 什么是Python
后端·python
golang学习记44 分钟前
Claude Code2.0.12发布插件系统:构建可扩展的 AI 编程工作流
人工智能
程序猿小D1 小时前
【完整源码+数据集+部署教程】遥感图像道路检测分割系统源码和数据集:改进yolo11-CARAFE
python·yolo·计算机视觉·目标跟踪·数据集·yolo11·遥感图像道路检测分割系统
minhuan1 小时前
构建AI智能体:六十二、金融风控系统:基于信息熵和KL散度的异常交易检测
人工智能·信息熵·kl散度·信息论·金融风控
zy_destiny1 小时前
【工业场景】用YOLOv8实现人员打电话识别
人工智能·python·深度学习·yolo·机器学习·计算机视觉·目标跟踪