欢迎关注 “小白玩转Python”,发现更多 “有趣”
引言
如果你刚刚开始学习 PyTorch,并想学习如何做一些基本的图像分类,那么这篇文章你一定不要错过哦~本文将通过组织自己的训练数据,使用预训练的神经网络来训练自己的模型,最终实现自己的图像分类!
组织训练数据集
PyTorch中的数据集是按照文件夹组织,并且每类都有一个文件夹。大多数教程都是在使用训练集和验证集的时候来进一步组织数据。但这是非常麻烦的,必须从每个类中选择一定数量的图像,然后将它们从训练集文件夹移动到验证集文件夹,这个过程又复杂又可能产生很多问题。
因此本文将引入一个更好的方法来将数据集分割成训练集和验证集。首先,让我们从导入模块开始:
%matplotlib inline
%config InlineBackend.figure_format = 'retina'
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
from torchvision import datasets, transforms, models
接下来定义训练数据集和验证数据集加载器,使用 SubsetRandomSampler 进行分割:
data_dir = '/data/train'
def load_split_train_test(datadir, valid_size = .2):
train_transforms = transforms.Compose([transforms.Resize(224),
transforms.ToTensor(),
])
test_transforms = transforms.Compose([transforms.Resize(224),
transforms.ToTensor(),
])
train_data = datasets.ImageFolder(datadir,
transform=train_transforms)
test_data = datasets.ImageFolder(datadir,
transform=test_transforms)
num_train = len(train_data)
indices = list(range(num_train))
split = int(np.floor(valid_size * num_train))
np.random.shuffle(indices)
from torch.utils.data.sampler import SubsetRandomSampler
train_idx, test_idx = indices[split:], indices[:split]
train_sampler = SubsetRandomSampler(train_idx)
test_sampler = SubsetRandomSampler(test_idx)
trainloader = torch.utils.data.DataLoader(train_data,
sampler=train_sampler, batch_size=64)
testloader = torch.utils.data.DataLoader(test_data,
sampler=test_sampler, batch_size=64)
return trainloader, testloader
trainloader, testloader = load_split_train_test(data_dir, .2)
print(trainloader.dataset.classes)
接下来检测一下我们的机器是否拥有GPU用于训练(使用GPU进行训练速度至少要快十倍哦~),然后加载一个预训练模型(在这里我们使用的是ResNet50):
device = torch.device("cuda" if torch.cuda.is_available()
else "cpu")
model = models.resnet50(pretrained=True)
print(model)
训练数据集
现在我们进入深层神经网络的有趣部分。首先必须冻结预训练的层,这样就不会在训练期间使用他们。然后重新定义完全连接的层。最后创建标准损失函数,并选择了一个优化器和学习率。
for param in model.parameters():
param.requires_grad = False
model.fc = nn.Sequential(nn.Linear(2048, 512),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(512, 10),
nn.LogSoftmax(dim=1))
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.fc.parameters(), lr=0.003)
model.to(device)
最后,让我们来训练我们的模型!在本例中只有一个epoch,但是在大多数情况下会需要更多epoch。从代码中可以很直观地看到这个基本过程: 加载图像并执行前馈循环。然后计算损失函数,并利用优化器将梯度下降应用到反向传播中。
PyTorch就是这么简单方便,下面是循环迭代过程:
epochs = 1
steps = 0
running_loss = 0
print_every = 10
train_losses, test_losses = [], []
for epoch in range(epochs):
for inputs, labels in trainloader:
steps += 1
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
logps = model.forward(inputs)
loss = criterion(logps, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if steps % print_every == 0:
test_loss = 0
accuracy = 0
model.eval()
with torch.no_grad():
for inputs, labels in testloader:
inputs, labels = inputs.to(device),
labels.to(device)
logps = model.forward(inputs)
batch_loss = criterion(logps, labels)
test_loss += batch_loss.item()
ps = torch.exp(logps)
top_p, top_class = ps.topk(1, dim=1)
equals =
top_class == labels.view(*top_class.shape)
accuracy +=
torch.mean(equals.type(torch.FloatTensor)).item()
train_losses.append(running_loss/len(trainloader))
test_losses.append(test_loss/len(testloader))
print(f"Epoch {epoch+1}/{epochs}.. "
f"Train loss: {running_loss/print_every:.3f}.. "
f"Test loss: {test_loss/len(testloader):.3f}.. "
f"Test accuracy: {accuracy/len(testloader):.3f}")
running_loss = 0
model.train()
torch.save(model, 'aerialmodel.pth')
然后等待几分钟(或者更长时间,这取决于数据集的大小和epoch的数量)之后,训练就结束了,将模型保存下来以便以后进行预测!
现在我们还可以绘制训练集和验证集的损失:
plt.plot(train_losses, label='Training loss')
plt.plot(test_losses, label='Validation loss')
plt.legend(frameon=False)
plt.show()
实现自己的图像分类
预测特定图像类别的功能非常简单。
def predict_image(image):
image_tensor = test_transforms(image).float()
image_tensor = image_tensor.unsqueeze_(0)
input = Variable(image_tensor)
input = input.to(device)
output = model(input)
index = output.data.cpu().numpy().argmax()
return index
现在为了更容易测试,我还创建了一个函数,可以从数据集文件夹中随机选取一些图片:
def get_random_images(num):
data = datasets.ImageFolder(data_dir, transform=test_transforms)
classes = data.classes
indices = list(range(len(data)))
np.random.shuffle(indices)
idx = indices[:num]
from torch.utils.data.sampler import SubsetRandomSampler
sampler = SubsetRandomSampler(idx)
loader = torch.utils.data.DataLoader(data,
sampler=sampler, batch_size=num)
dataiter = iter(loader)
images, labels = dataiter.next()
return images, labels
最后为了演示预测函数,我们使用随机的图像样本,对其进行预测并显示结果:
to_pil = transforms.ToPILImage()
images, labels = get_random_images(5)
fig=plt.figure(figsize=(10,10))
for ii in range(len(images)):
image = to_pil(images[ii])
index = predict_image(image)
sub = fig.add_subplot(1, len(images), ii+1)
res = int(labels[ii]) == index
sub.set_title(str(classes[index]) + ":" + str(res))
plt.axis('off')
plt.imshow(image)
plt.show()
这是在Google地图图块上进行分类的一个示例。标签是预测的类别,并且显示它是否是正确的预测。
结束语
在本文中我们学习了如何通过PyTorch实现自己的图像分类,是不是觉得很有趣~当然PyTorch的可以做的事情还有更多,欢迎大家继续探索下去!
· END ·
HAPPY LIFE