在本教程中，将介绍如何使用Torchvision中的Faster R-CNN模型对自定义数据集进行目标检测微调。我们将使用主流的Penn-Fudan数据集，这个数据集具有170个图像和345个行人实例，用于展示如何在一个具体的场景中进行行人检测。由于Penn-Fudan数据集比较小，所以我们将从COCO数据集中使用预先训练的模型来进行微调。
具体的，这个教程包含以下内容：
- 加载数据集
- 定义模型
- 定义损失函数
- 定义优化器
- 训练和验证模型
相信通过仔细学习，你会构建自己的深度学习目标检测模型！
如果对你有帮助还请多多关注，时刻分享人工智能相关知识！！
1. 加载数据集
我们将使用Penn-Fudan人工智能数据集，该数据集提供了170张包含行人实例的图像。我们已经将该数据集分成了训练集和验证集，其中训练集包含100张图像，验证集包含70张图像。这是一个二元分类问题：对于每个图像，我们需要预测是否存在行人实例。
Penn-Fudan数据集的文件夹结构如下：
PennFudanPed/
PedMasks/
FudanPed00001_mask.png
FudanPed00002_mask.png
FudanPed00003_mask.png
FudanPed00004_mask.png
...
PNGImages/
FudanPed00001.png
FudanPed00002.png
FudanPed00003.png
FudanPed00004.png
...
Ped_annotations/
FudanPed00001.txt
FudanPed00002.txt
FudanPed00003.txt
FudanPed00004.txt
...
我们将使用torchvision提供的PennFudanDataset类来加载数据集。PennFudanDataset类继承自torch.utils.data.Dataset类，因此我们可以使用torchvision提供的数据加载器来加载数据集。
数据集下载可以联系我---
下面是加载数据集的代码：
相信做过很多项目的你，很熟悉torch加载数据的风格！！
import os
import numpy as np
import torch
from PIL import Image
class PennFudanDataset(torch.utils.data.Dataset):
def __init__(self, root, transforms=None):
self.root = root
self.transforms = transforms
# load all image files, sorting them to
# ensure that they are aligned
self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages"))))
self.masks = list(sorted(os.listdir(os.path.join(root, "PedMasks"))))
def __getitem__(self, idx):
# load images and masks
img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])
mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
img = Image.open(img_path).convert("RGB")
# note that we haven't converted the mask to RGB,
# because each color corresponds to a different instance
# with 0 being background
mask = Image.open(mask_path)
# convert the PIL Image into a numpy array
mask = np.array(mask)
# instances are encoded as different colors
obj_ids = np.unique(mask)
# first id is the background, so remove it
obj_ids = obj_ids[1:]
# split the color-encoded mask into a set
# of binary masks
masks = mask == obj_ids[:, None, None]
# get bounding box coordinates for each mask
num_objs = len(obj_ids)
boxes = []
for i in range(num_objs):
pos = np.where(masks[i])
xmin = np.min(pos[1])
xmax = np.max(pos[1])
ymin = np.min(pos[0])
ymax = np.max(pos[0])
boxes.append([xmin, ymin, xmax, ymax])
# convert everything into a torch.Tensor
boxes = torch.as_tensor(boxes, dtype=torch.float32)
# there is only one class
labels = torch.ones((num_objs,), dtype=torch.int64)
masks = torch.as_tensor(masks, dtype=torch.uint8)
image_id = torch.tensor([idx])
数据加载完毕，我们开始构建模型。
2. 定义模型
我们将使用Torchvision中提供的高精度目标检测模型Faster R-CNN模型进行微调。Faster R-CNN是一种基于区域提取的目标检测方法，它通过先提取候选区域，再对这些候选区域进行分类和回归，从而检测图像中的目标。 在本教程中，我们将使用预先训练的Faster R-CNN模型进行微调，而不是从头开始训练一个新模型。我们将使用COCO数据集提供的预先训练的Faster R-CNN模型，该模型已经对COCO数据集进行了训练，可以很好地提取图像中的目标区域**。 这也是一个经典的两阶段目标检测器！**
下面是定义模型的代码：
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
# load a pre-trained model for classification and return
# only the features
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# FasterRCNN needs to know the number of
# output channels in a backbone. For mobilenet_v2, it's 1280
# so we need to add it here
backbone.out_channels = 1280
# let's make the RPN generate 5 x 3 anchors per spatial
# location, with 5 different sizes and 3 different aspect
# ratios. We have a Tuple[Tuple[int]] because each feature
# map could potentially have different sizes and
# aspect ratios
anchor_generator = torchvision.models.detection.rpn.AnchorGenerator(
sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),)
)
# let's define what are the feature maps that we will
# use to perform the region of interest cropping, as well as
# the size of the crop after rescaling.
# if your backbone returns a Tensor, featmap_names is expected to
# be [0]. More generally, the backbone should return an
# OrderedDict[Tensor], and in featmap_names you can choose which
# feature maps to use.
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=['0'],
output_size=7,
sampling_ratio=2
)
# put the pieces together inside a FasterRCNN model
model = torchvision.models.detection.FasterRCNN(
backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler
)
# replace the classifier with a new one, that has
# num_classes which is user-defined
num_classes = 2
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
在上面的代码中，我们首先加载了一个预训练的轻量级网络模型MobileNetV2模型，并取出了该模型的特征提取器部分。然后，我们将特征提取器的输出通道数设置为1280，以使其与Faster R-CNN模型兼容。
接下来，我们定义了一个锚点生成器和一个区域兴趣池，它们将在Faster R-CNN模型中用于生成候选区域和对这些候选区域进行分类和回归。最后，我们将Faster R-CNN模型的分类器替换为一个新的分类器，该分类器具有用户定义的类别数。
3. 定义损失函数
我们将使用Faster R-CNN模型默认的损失函数，即Multi-task loss多任务损失。Multi-task loss由两个部分组成：分类损失和回归损失。分类损失用于衡量模型对候选区域进行分类的准确性，回归损失用于衡量模型对候选区域进行回归的准确性。
下面是定义损失函数的代码：
import torch.nn as nn
import torch.optim as optim
# move model to the right device
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)
4. 定义优化器
我们将使用随机梯度下降（SGD）优化器（相信熟悉深度学习的你，一定耳熟能详了这个算法）来训练模型。我们还将使用学习率调度程序来自适应地调整学习率，以提高模型的性能。
下面是定义优化器和学习率调度器的代码：
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,
momentum=0.9, weight_decay=0.0005)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
在上面的代码中，我们首先获取所有需要进行梯度更新的模型参数，并将它们传递给SGD优化器。我们还指定了学习率、动量和权重衰减等参数。
然后，我们定义了一个学习率调度器，它将在每个epoch结束时自适应地调整学习率。我们将使用StepLR调度程序，每3个epoch将学习率降低10倍。
5. 训练和验证模型
我们将使用PyTorch提供的训练和验证循环来训练和验证我们的模型。
下面是训练和验证模型的代码：
from engine import train_one_epoch, evaluate
# let's train it for 10 epochs
num_epochs = 10
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
evaluate(model, data_loader_test, device=device)
在上面的代码中，我们首先指定了要训练的epoch数。然后，我们使用for循环迭代训练和验证模型。
在每个epoch中，我们使用train_one_epoch函数训练模型，该函数将模型、优化器、数据加载器、设备和当前epoch作为输入，并打印每10个batch的训练损失和时间。
在每个epoch结束时，我们使用学习率调度器自适应地调整学习率。然后，我们使用evaluate函数在验证集上评估模型，该函数将模型、数据加载器和设备作为输入，并打印出平均精度和平均损失。
最后，我们完成了目标检测微调的教程。您可以将本教程中的代码与自己的数据集一起使用，以进行目标检测微调。
6. 保存和加载模型
在训练和验证模型之后，我们可能希望保存模型以备将来使用。我们可以使用PyTorch提供的save函数将模型保存到磁盘上。
下面是保存模型的代码：# save the model to disk
torch.save(model.state_dict(), "model.pth")
在上面的代码中，我们使用state_dict函数获取模型的状态字典，并使用torch.save函数将其保存到名为"model.pth"的文件中。
如果我们希望在将来加载模型，我们可以使用load_state_dict函数从磁盘上的文件中加载模型的状态字典。
下面是加载模型的代码：
# create an instance of the model and load the state dict
model = get_instance_segmentation_model(num_classes)
model.load_state_dict(torch.load("model.pth"))
在上面的代码中，我们首先使用get_instance_segmentation_model函数创建一个模型实例。然后，我们使用load_state_dict函数从名为"model.pth"的文件中加载模型的状态字典。
7. 结论
本教程介绍了如何使用PyTorch进行目标检测微调。我们首先了解了目标检测的基本概念和检测器的架构。然后，我们介绍了如何使用PyTorch构建目标检测模型，并使用预训练的模型进行微调。最后，我们介绍了如何保存和加载模型。
PyTorch提供了强大的工具和库，使得目标检测微调变得非常容易。希望这个教程能够帮助您进一步探索目标检测领域，并构建更准确和高效的检测器。
扫一扫
244
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
