对于这个任务,我们需要使用深度学习模型来对宠物猫进行分类。由于数据集比较大,我们可以使用预训练好的模型来进行迁移学习。这里我们选择使用ResNet模型来进行分类。

首先,我们需要导入必要的库和模块:

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms

然后,我们需要定义一些超参数,包括批大小、学习率、训练轮数等:

batch_size = 64
learning_rate = 0.001
num_epochs = 10

接下来,我们需要定义数据预处理的方式。这里我们使用了常用的图像预处理方式,包括随机裁剪、随机翻转、归一化等:

train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

然后,我们需要定义数据加载器。这里我们使用PyTorch自带的ImageFolder类来加载数据:

train_dataset = ImageFolder('kaggle/input/cat_test/train', transform=train_transform)
test_dataset = ImageFolder('kaggle/input/cat_test/test', transform=test_transform)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

接下来,我们需要定义ResNet模型。这里我们使用ResNet18作为基础模型,并替换掉最后一层全连接层:

class ResNet(nn.Module):
    def __init__(self):
        super(ResNet, self).__init__()
        self.resnet = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True)
        num_features = self.resnet.fc.in_features
        self.resnet.fc = nn.Linear(num_features, 12)

    def forward(self, x):
        x = self.resnet(x)
        return x

接下来,我们需要定义损失函数和优化器:

model = ResNet().cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

然后,我们就可以开始训练模型了:

for epoch in range(num_epochs):
    train_loss, train_acc = 0.0, 0.0
    model.train()
    for images, labels in train_loader:
        images = images.cuda()
        labels = labels.cuda()

        optimizer.zero_grad()

        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item() * images.size(0)
        _, preds = torch.max(outputs, 1)
        train_acc += torch.sum(preds == labels.data)

    train_loss = train_loss / len(train_loader.dataset)
    train_acc = train_acc.double() / len(train_loader.dataset)

    print('Epoch {}: Train Loss: {:.4f}, Train Acc: {:.4f}'.format(
        epoch+1, train_loss, train_acc))

    model.eval()
    test_loss, test_acc = 0.0, 0.0
    for images, labels in test_loader:
        images = images.cuda()
        labels = labels.cuda()

        with torch.no_grad():
            outputs = model(images)
            loss = criterion(outputs, labels)

            test_loss += loss.item() * images.size(0)
            _, preds = torch.max(outputs, 1)
            test_acc += torch.sum(preds == labels.data)

    test_loss = test_loss / len(test_loader.dataset)
    test_acc = test_acc.double() / len(test_loader.dataset)

    print('Epoch {}: Test Loss: {:.4f}, Test Acc: {:.4f}'.format(
        epoch+1, test_loss, test_acc))

最后,我们可以使用训练好的模型对测试集进行分类,并输出分类结果:

model.eval()
preds = []
with torch.no_grad():
    for images, _ in test_loader:
        images = images.cuda()
        outputs = model(images)
        _, pred = torch.max(outputs, 1)
        preds += pred.tolist()

print(preds)

完整代码如下:

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms

batch_size = 64
learning_rate = 0.001
num_epochs = 10

train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

test_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

train_dataset = ImageFolder('kaggle/input/cat_test/train', transform=train_transform)
test_dataset = ImageFolder('kaggle/input/cat_test/test', transform=test_transform)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

class ResNet(nn.Module):
    def __init__(self):
        super(ResNet, self).__init__()
        self.resnet = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True)
        num_features = self.resnet.fc.in_features
        self.resnet.fc = nn.Linear(num_features, 12)

    def forward(self, x):
        x = self.resnet(x)
        return x

model = ResNet().cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    train_loss, train_acc = 0.0, 0.0
    model.train()
    for images, labels in train_loader:
        images = images.cuda()
        labels = labels.cuda()

        optimizer.zero_grad()

        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item() * images.size(0)
        _, preds = torch.max(outputs, 1)
        train_acc += torch.sum(preds == labels.data)

    train_loss = train_loss / len(train_loader.dataset)
    train_acc = train_acc.double() / len(train_loader.dataset)

    print('Epoch {}: Train Loss: {:.4f}, Train Acc: {:.4f}'.format(
        epoch+1, train_loss, train_acc))

    model.eval()
    test_loss, test_acc = 0.0, 0.0
    for images, labels in test_loader:
        images = images.cuda()
        labels = labels.cuda()

        with torch.no_grad():
            outputs = model(images)
            loss = criterion(outputs, labels)

            test_loss += loss.item() * images.size(0)
            _, preds = torch.max(outputs, 1)
            test_acc += torch.sum(preds == labels.data)

    test_loss = test_loss / len(test_loader.dataset)
    test_acc = test_acc.double() / len(test_loader.dataset)

    print('Epoch {}: Test Loss: {:.4f}, Test Acc: {:.4f}'.format(
        epoch+1, test_loss, test_acc))

model.eval()
preds = []
with torch.no_grad():
    for images, _ in test_loader:
        images = images.cuda()
        outputs = model(images)
        _, pred = torch.max(outputs, 1)
        preds += pred.tolist()

print(preds)
宠物猫分类:使用ResNet模型进行12种猫的分类

原文地址: https://www.cveoy.top/t/topic/oD3J 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录