对我来说，图像定位是一个有趣的应用，因为它正好介于图像分类和物体检测之间。这是使用PyTorch的物体定位系列的第二部分，如果你没有在这里，请查看前一部分。
请务必关注Gradient上的IPython笔记本，并叉开你自己的版本来尝试一下
数据集的可视化
在进入本教程的机器学习部分之前，让我们用边界框对数据集进行可视化。在这里，我们可以看到如何通过与图像大小相乘来检索坐标。我们使用OpenCV来显示图像。这将打印出20张图像的样本及其边界框。
# Create a Matplotlib figure
plt.figure(figsize=(20,20));
# Generate a random sample of images each time the cell is run
random_range = random.sample(range(1, len(img_list)), 20)
for itr, i in enumerate(random_range, 1):
# Bounding box of each image
a1, b1, a2, b2 = boxes[i];
img_size = 256
# Rescaling the boundig box values to match the image size
x1 = a1 * img_size
x2 = a2 * img_size
y1 = b1 * img_size
y2 = b2 * img_size
# The image to visualize
image = img_list[i]
# Draw bounding boxes on the image
cv2.rectangle(image, (int(x1),int(y1)),
(int(x2),int(y2)),
(0,255,0),
3);
# Clip the values to 0-1 and draw the sample of images
img = np.clip(img_list[i], 0, 1)
plt.subplot(4, 5, itr);
plt.imshow(img);
plt.axis('off');
数据集的可视化
数据集的分割
我们已经得到了Img_list、标签和盒子中的数据集，现在我们必须在跳转到DataLoaders之前分割数据集。像往常一样，我们将使用sklearn库中的train_test_split方法完成这项任务。
# Split the data of images, labels and their annotations
train_images, val_images, train_labels, \
val_labels, train_boxes, val_boxes = train_test_split( np.array(img_list),
np.array(labels), np.array(boxes), test_size = 0.2,
random_state = 43)
print('Training Images Count: {}, Validation Images Count: {}'.format(
len(train_images), len(val_images) ))
现在我们已经通过可视化快速浏览了我们的数据集，并完成了数据集的分割。让我们继续为我们的数据集建立自定义的PyTorch DataLoaders，这些数据集目前散落在各个变量中。
Pytorch中的自定义DataLoaders
DataLoaders，顾名思义，会返回一个对象，在我们训练模型时处理整个数据供应系统。它在创建对象时提供了shuffle等功能，它有一个'getitem'方法来处理每个迭代中的数据输入，所有这些东西让你以你希望的方式设计一切，而不会使训练部分的代码变得混乱。这让你可以更专注于其他优化。让我们从PyTorch部分的导入开始。
from PIL import Image
import torch
import torchvision
from torchvision.transforms import ToTensor
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import os
import pickle
import random
import time
如果可以的话，重要的事情之一是使用GPU来训练ML模型，特别是当目标巨大时。如果在Paperspace Gradient中运行，请选择一台有GPU的机器。
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device
>>>device(type='cuda')
上面的输出表明你有一个GPU，设备可以用来将数据和模型转换到GPU上进行利用。继续前进到DataLoaders。
class Dataset():
def __init__(self, train_images, train_labels, train_boxes):
self.images = torch.permute(torch.from_numpy(train_images),(0,3,1,2)).float()
self.labels = torch.from_numpy(train_labels).type(torch.LongTensor)
self.boxes = torch.from_numpy(train_boxes).float()
def __len__(self):
return len(self.labels)
# To return x,y values in each iteration over dataloader as batches.
def __getitem__(self, idx):
return (self.images[idx],
self.labels[idx],
self.boxes[idx])
# Inheriting from Dataset class
class ValDataset(Dataset):
def __init__(self, val_images, val_labels, val_boxes):
self.images = torch.permute(torch.from_numpy(val_images),(0,3,1,2)).float()
self.labels = torch.from_numpy(val_labels).type(torch.LongTensor)
self.boxes = torch.from_numpy(val_boxes).float()
在这里，我们将创建Dataset类，首先加载图像、标签和盒子的坐标，将其缩放到[0-1]的范围内到类变量。然后我们使用 "getitem "来设计每个迭代中的加载器输出。同样地，我们将建立ValDataset（验证数据集）DataLoader类。由于数据的结构和性质是相同的，我们将继承上述类。
现在我们有了DataLoaders的类，让我们马上从各自的类中创建数据加载器对象。
dataset = Dataset(train_images, train_labels, train_boxes)
valdataset = ValDataset(val_images, val_labels, val_boxes)
现在我们已经完成了准备数据的过程，让我们转到教程的机器学习模型部分。我们将使用一套相当简单的卷积神经网络来进行对象定位的实现，这将有助于我们根据自己的理解来调整这个概念。我们可以对这个架构的各个方面进行实验，包括使用像Alexnet这样的预训练模型。我鼓励你学习更多关于优化器、损失函数和超调模型的知识，为未来的项目掌握架构设计的技能。
模型架构
因此，为了理解我们必须如何设计架构，我们首先要理解输入和输出。这里的输入是一批图像，所以它将被风格化为(BS, C, H, W) 。BS 是批次大小，然后是通道、高度和重量。这个顺序很重要，因为在PyTorch中，图像是这样存储的。在TensorFlow中，每个图像的顺序是（H，W，C）。
说到输出，我们有两个输出，就像我们在之前的博客中开始讨论的那样。首先是你的分类输出，它的大小为(1, N) ，N 是类的数量。第二个输出是大小为(1, 4) ，即xmin、ymin、xmax和ymax，但范围为(0,1) 。这有助于你以后根据图像大小来调整坐标。因此，输出将是(CLF, BOX) ，第一个是上面讨论的分类，其他是坐标。
现在，从机器学习部分来解释，CLF ，通过在输出上使用argmax来获得类指数，我们可以在最后添加一个softmax激活函数来获得概率输出。BOX 输出将在通过sigmoid激活函数后产生，这使得范围(0,1) 。也就是说，其他一切都与传统的图像分类模型类似。让我们看一下代码。
class Network(nn.Module):
def __init__(self):
super(Network, self).__init__()
# CNNs for rgb images
self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=5)
self.conv2 = nn.Conv2d(in_channels=6, out_channels=12, kernel_size=5)
self.conv3 = nn.Conv2d(in_channels=12, out_channels=24, kernel_size=5)
self.conv4 = nn.Conv2d(in_channels=24, out_channels=48, kernel_size=5)
self.conv5 = nn.Conv2d(in_channels=48, out_channels=192, kernel_size=5)
# Connecting CNN outputs with Fully Connected layers for classification
self.class_fc1 = nn.Linear(in_features=1728, out_features=240)
self.class_fc2 = nn.Linear(in_features=240, out_features=120)
self.class_out = nn.Linear(in_features=120, out_features=2)
# Connecting CNN outputs with Fully Connected layers for bounding box
self.box_fc1 = nn.Linear(in_features=1728, out_features=240)
self.box_fc2 = nn.Linear(in_features=240, out_features=120)
self.box_out = nn.Linear(in_features=120, out_features=4)
def forward(self, t):
t = self.conv1(t)
t = F.relu(t)
t = F.max_pool2d(t, kernel_size=2, stride=2)
t = self.conv2(t)
t = F.relu(t)
t = F.max_pool2d(t, kernel_size=2, stride=2)
t = self.conv3(t)
t = F.relu(t)
t = F.max_pool2d(t, kernel_size=2, stride=2)
t = self.conv4(t)
t = F.relu(t)
t = F.max_pool2d(t, kernel_size=2, stride=2)
t = self.conv5(t)
t = F.relu(t)
t = F.avg_pool2d(t, kernel_size=4, stride=2)
t = torch.flatten(t,start_dim=1)
class_t = self.class_fc1(t)
class_t = F.relu(class_t)
class_t = self.class_fc2(class_t)
class_t = F.relu(class_t)
class_t = F.softmax(self.class_out(class_t),dim=1)
box_t = self.box_fc1(t)
box_t = F.relu(box_t)
box_t = self.box_fc2(box_t)
box_t = F.relu(box_t)
box_t = self.box_out(box_t)
box_t = F.sigmoid(box_t)
return [class_t,box_t]
让我们接下来实例化这个模型，并让它使用GPU（如果有）。这确实可以加速训练过程，特别是对于像图像定位这样的巨大目标。在这之后，我们还将看到模型的总结。
model = Network()
model = model.to(device)
model
我们可以在以后使用ONNX格式和netron网站在训练后将模型可视化，以评估复杂的细节，如权重和更多。现在，我们可以写一个小函数来计算每批训练的正确预测数。
def get_num_correct(preds, labels):
return torch.round(preds).argmax(dim=1).eq(labels).sum().item()
现在我们将为训练和验证数据集创建数据加载器，以便在训练时输入成批的图像。
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=32, shuffle=True)
valdataloader = torch.utils.data.DataLoader(
valdataset, batch_size=32, shuffle=True)
在这里，我们可以在Dataloader方法中设置更多的参数，如num_workers，这可以帮助我们改善加载时间，其中2通常是最佳值，可以实现数据加载过程的管道化。现在我们要创建训练函数，这是每个ML项目中最重要的部分。我们将保持这个过程的简单性，让大家了解训练的每个部分。我们将专注于通过这个过程提高验证的准确性。
把这个项目带入生活
模型训练和验证
def train(model):
# Defining the optimizer
optimizer = optim.SGD(model.parameters(),lr = 0.1)
num_of_epochs = 30
epochs = []
losses = []
# Creating a directory for storing models
os.mkdir('models')
for epoch in range(num_of_epochs):
tot_loss = 0
tot_correct = 0
train_start = time.time()
model.train()
for batch, (x, y, z) in enumerate(dataloader):
# Converting data from cpu to GPU if available to improve speed
x,y,z = x.to(device),y.to(device),z.to(device)
# Sets the gradients of all optimized tensors to zero
optimizer.zero_grad()
[y_pred,z_pred]= model(x)
# Compute loss (here CrossEntropyLoss)
class_loss = F.cross_entropy(y_pred, y)
box_loss = F.mse_loss(z_pred, z)
(box_loss + class_loss).backward()
# class_loss.backward()
optimizer.step()
print("Train batch:", batch+1, " epoch: ", epoch, " ",
(time.time()-train_start)/60, end='\r')
model.eval()
for batch, (x, y,z) in enumerate(valdataloader):
# Converting data from cpu to GPU if available to improve speed
x,y,z = x.to(device),y.to(device),z.to(device)
# Sets the gradients of all optimized tensors to zero
optimizer.zero_grad()
with torch.no_grad():
[y_pred,z_pred]= model(x)
# Compute loss (here CrossEntropyLoss)
class_loss = F.cross_entropy(y_pred, y)
box_loss = F.mse_loss(z_pred, z)
# Compute loss (here CrossEntropyLoss)
tot_loss += (class_loss.item() + box_loss.item())
tot_correct += get_num_correct(y_pred, y)
print("Test batch:", batch+1, " epoch: ", epoch, " ",
(time.time()-train_start)/60, end='\r')
epochs.append(epoch)
losses.append(tot_loss)
print("Epoch", epoch, "Accuracy", (tot_correct)/2.4, "loss:",
tot_loss, " time: ", (time.time()-train_start)/60, " mins")
torch.save(model.state_dict(), "model_ep"+str(epoch+1)+".pth")
让我们详细了解一下。我们将使用的优化器是随机梯度下降，但如果你愿意，你可以尝试其他的优化技术，如亚当。
我们将单独处理训练和验证问题。为了在训练期间跟踪时间，我们可以使用train_start 变量。事实证明，当在付费的GPU实例上工作时，这对计算训练的时间相当有用，这将有助于你在多次重新训练之前提前计划。
默认情况下，PyTorch模型被设置为训练（self.training = True）。我们将讨论为什么我们要切换状态，以及当我们到达代码的验证部分时的效果。
从我们创建的数据加载器中，我们访问由x、y和z组成的每批数据，它们分别是图像、标签和边界框。然后，我们将它们中的每一个转换到我们的首选设备上，也就是转换到GPU上，如果它是可用的。优化器处理深度学习中的反向传播，所以在训练之前，我们将每一批的梯度设置为零，这是通过optimizer.zero_grad() 完成的*。*
在输入(x)到模型后，我们进入损失计算阶段。这是一个重要的部分，因为这里涉及两种类型的损失。用于分类问题的交叉熵损失，以及用于寻找边界盒坐标的回归部分的平均平方误差。如果我们观察这里，我们可以看到我们是如何为反向传播发送损失之和，而不是单独处理。
之后，代码的评估部分与训练部分类似，只是我们没有做反向传播。一旦我们进入验证部分，我们就用*model.eval()*将模型状态切换为评价。正如我们前面所讨论的，模型默认处于训练状态，而这种切换非常重要。一些层在训练/评估期间有不同的行为（如BatchNorm，Dropout），这反过来又会影响整体性能。这是一个有趣的观点，因为当模型处于训练状态时，BatchNorm使用每批统计，Dropout层也被激活。但是当模型处于评估（推理）模式时，BatchNorm层使用运行统计，而Dropout层则被停用。这里需要注意的是，这些函数调用都不是向前/向后运行的。它们告诉模型在运行时如何行动。这一点很重要，因为一些模块（层）被设计成在训练和推理期间有不同的行为，如果在错误的模式下运行，模型会产生意想不到的结果。在这里，我们用之前写的函数来计算准确率，找出正确预测的数量。
在这个例子中，我们将训练30个epochs，但这是任意的，我们可以自由尝试更长或更短的epoch值。
关于训练过程的最后一件事，我们必须选择的是如何保存模型。在PyTorch中流行的是，我们可以通过两种方式保存模型。一种是我们将整个模式原封不动地保存下来，包括权重、结构等；如果我们确信生产的PyTorch版本与开发的版本相匹配，那么这种保存方式是完美的。在某些情况下，当你有500MB左右的大小限制时，以这种方式保存的PyTorch模型对于部署来说可能过于沉重。另一种保存模型的方法是通过创建state_dict()， 它只将模型的可学习参数的状态保存为python字典。这使得它在形式上更加模块化，因为在生产中，要复活模型，你必须提供模型的state_dict()和模型的架构*（在我们的例子中是网络*）。这几乎就像提供肌肉和骨架来复活一个生命，这很方便和灵活，因为我们可以决定PyTorch的版本，如果需要，可以使用PyTorch的cpu专用版本，这将减少整个包的大小，这种方法非常有用。我鼓励你阅读更多关于 state_dict()的内容，以扩大你的理解。也就是说，如果没有内存存储限制，保存所有的模型数据是最安全的方法。否则选择像早期停止这样的技术将帮助你获得最好的模型，而不会错过。
一旦你理解了训练功能，让我们继续开始训练。
train(model)
这里需要理解的一件事是，如果你想重新训练，你必须重新初始化模型。没有这一步的重新训练将只是对现有模型参数的补充。因此，处理这个潜在问题的一个好方法是在训练前初始化模型，再次重复同样的工作。
验证的准确性取决于很多东西，你可能不会在第一次尝试时就得到一个好的验证。我提供的测试实现的数据集是一个最小的数据集，你必须有一个更详细和更普遍的数据集来对你的模型进行真正的测试。其他一些改进模型的方法是通过数据增强技术等。我已经写了一篇 关于提高准确性的详细文章，如果你的模型仍然需要改进，请阅读它。
模型测试/推理
当涉及到图像定位时，预测不只是从模型中获得一个输出。我们还必须处理边界框坐标，以生成一个实际的边界框来可视化结果，这甚至可能有助于生产。本节将有3个主要部分，预测脚本、前处理和后处理脚本，所以让我们跳入其中。
前处理
在进行机器学习制作时，重要的是你提供给模型的数据的预处理方式要与你在训练时提供的数据完全相同。调整大小是任何图像相关模型在推理阶段最常见的预处理之一。 这里我们的图像大小为256。
def preprocess(img, image_size = 256):
image = cv2.resize(img, (image_size, image_size))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = image.astype("float") / 255.0
# Expand dimensions as predict expect image in batches
image = np.expand_dims(image, axis=0)
return image
后期处理
一旦我们得到输出，它们将以[CLF, BOX] ，以更简单的方式，我们将不得不与边界盒的值一起工作，以创造具有可视化的结果。边界框输入被缩放到[0,1] ，而且，由于我们在最后使用了sigmoid激活函数，我们的预测值也在[0,1] ，范围内。我们必须对它们进行重新缩放，以得到xmin、ymin等。为此，我们只需将这些值与图像大小（此处为256）相乘。
def postprocess(image, results):
# Split the results into class probabilities and box coordinates
[class_probs, bounding_box] = results
# First let's get the class label
# The index of class with the highest confidence is our target class
class_index = torch.argmax(class_probs)
# Use this index to get the class name.
class_label = num_to_labels[class_index]
# Now you can extract the bounding box too.
# Get the height and width of the actual image
h, w = 256,256
# Extract the Coordinates
x1, y1, x2, y2 = bounding_box[0]
# # Convert the coordinates from relative (i.e. 0-1) to actual values
x1 = int(w * x1)
x2 = int(w * x2)
y1 = int(h * y1)
y2 = int(h * y2)
# return the lable and coordinates
return class_label, (x1,y1,x2,y2),torch.max(class_probs)*100
现在我们有了前处理和后处理，让我们进入预测脚本。
预测脚本
在预测脚本中，我们首先从早期的网络中获得模型架构，然后将模型移到我们的首选设备上：梯度笔记本的GPU。一旦完成，我们就加载模型的state_dict()，正如我们之前在验证中讨论的那样。我们将模型设置为评估状态，并使用预处理函数使图像准备好被送入模型。[N,C,H,W] 然后我们可以使用PyTorch中的permute函数来重新制作图像数组，从[N,H,W,C] 。结果交给后处理函数，后处理函数给出实际的坐标和标签。最后，我们使用matplotlib来绘制带有边界框的图像。
# We will use this function to make prediction on images.
def predict(image, scale = 0.5):
model = Network()
model = model.to(device)
model.load_state_dict(torch.load("models/model_ep29.pth"))
model.eval()
# Reading Image
img = cv2.imread(image)
# # Before we can make a prediction we need to preprocess the image.
processed_image = preprocess(img)
result = model(torch.permute(torch.from_numpy(processed_image).float(),(0,3,1,2)).to(device))
# After postprocessing, we can easily use our results
label, (x1, y1, x2, y2), confidence = postprocess(image, result)
# Now annotate the image
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 100), 2)
cv2.putText(
img,
'{}, CONFIDENCE: {}'.format(label, confidence),
(30, int(35 * scale)),
cv2.FONT_HERSHEY_COMPLEX, scale,
(200, 55, 100),
2
)
# Show the Image with matplotlib
plt.figure(figsize=(10,10))
plt.imshow(img[:,:,::-1])
让我们继续在图像上进行推理：
image = '< IMAGE PATH >'
predict(image)
下面是输出结果，你可以看到我们是如何在图像上得到一个边界框的。
带有边界框的结果图像
ONNX和模型结构的可视化
正如我们前面提到的，我们可以使用netron来可视化我们的模型及其权重等。为此，我们可以将我们的模型转换成ONNX格式。ONNX格式是一种通用的机器学习模型存储格式，被Snap Lens Studio等平台所接受，它也是Gradient上可部署的框架之一。建立在TensorFlow和PyTorch上的模型可以转换为ONNX，其中两者实现了相同的形式和功能。以下是我们将如何做。在获得'.ONNX'文件后，你可以访问netron以获得更深入的模型可视化。
model = Network()
model = model.to(device)
model.load_state_dict(torch.load('models/model_ep29.pth'))
model.eval()
random_input = torch.randn(1, 3, 256, 256, dtype=torch.float32).to(device)
# you can add however many inputs your model or task requires
input_names = ["input"]
output_names = ["cls","loc"]
torch.onnx.export(model, random_input, 'localization_model.onnx', verbose=False,
input_names=input_names, output_names=output_names,
opset_version=11)
print("DONE")
结论
正如我们所讨论的那样，图像定位是一个有趣的研究领域，与图像分类问题相比，要想在这方面获得准确度是比较困难的。在这里，我们不仅要得到正确的分类，还要得到一个紧密而正确的边界盒。处理这两个问题会使模型变得复杂，所以改进数据集和架构会对你有帮助。我希望这篇关于实现图像定位的详细的两部分文章是有帮助的，我鼓励你阅读更多相同的内容。
