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()

部分测试结果如下

相关推荐
max50060032 分钟前
使用OmniAvatar-14B模型实现照片和文字生成视频的完整指南
图像处理·人工智能·深度学习·算法·音视频
不一样的故事12639 分钟前
学习Python是一个循序渐进的过程,结合系统学习、持续实践和项目驱动,
开发语言·python·学习
可触的未来,发芽的智生40 分钟前
追根索源-神经网络的灾难性遗忘原因
人工智能·神经网络·算法·机器学习·架构
CAE32041 分钟前
基于Ncode的新能源汽车电池包随机振动疲劳分析
人工智能·汽车·电池包·hypermesh·振动疲劳·optistruct
zzywxc78741 分钟前
自动化测试框架是软件测试的核心基础设施,通过预设规则和脚本自动执行测试用例,显著提高测试效率和覆盖率。
运维·人工智能·自动化·prompt·测试用例·流程图
尺度商业1 小时前
2025服贸会“海淀之夜”,点亮“科技”与“服务”底色
大数据·人工智能·科技
AWS官方合作商1 小时前
涂鸦智能携手亚马逊云科技,以全球基础设施与生成式AI加速万物智联时代到来
人工智能·科技·aws·亚马逊云科技
@HNUSTer1 小时前
Python数据可视化科技图表绘制系列教程(七)
python·数据可视化·科技论文·专业制图·科研图表
FunTester1 小时前
拥抱直觉与创造力:走进VibeCoding的新世界
人工智能·语言模型·编程·vibecoding
shizidushu1 小时前
How to work with merged cells in Excel with `openpyxl` in Python?
python·microsoft·excel·openpyxl