宠物猫分类:使用ResNet模型进行12种猫的分类
对于这个任务,我们需要使用深度学习模型来对宠物猫进行分类。由于数据集比较大,我们可以使用预训练好的模型来进行迁移学习。这里我们选择使用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)
原文地址: https://www.cveoy.top/t/topic/oD3J 著作权归作者所有。请勿转载和采集!