对于训练对象检测的脚本，实例分割和人员关键点检测要能够轻松支持添加新的自定义的数据，数据集应该从标准的类torch.utils.data.Dataset继承，并实现__len__和__getitem__
对于__getitem__应该返回：
数据集——>我在这里
下载后解压放在工作路径下，先查看一下数据集吧，结构如下：
（黑压压一大片，这没问题哈，是掩膜，哈哈哈哈哈~不要以为自己下载的不对！）
文件夹中每个文件的详细说明如下（readme中详细说明）：
在这里插入代码片
定义一个可以预测上述数据集的模型。本节，使用Mask R-CNN，它基于Faster R-CNN。Faster R-CNN是一个模型，可以预测图像中潜在对象的边界框和类别得分。（此处贴一个Faster R-CNN的详解——>我在这里）
Mask R-CNN在Faster R-CNN中添加了一个额外的分支，预测每个实例的分割蒙版
有两种常见情况可能需要修改torchvision modelzoo中的一个可用模型。
以下是读这两种情况的处理：
假设你想从一个在CoCo（其实它叫COCO ）上已预先训练过的模型开始，并希望为特定类进行微调，这是一种可行的方法：
#%%微调已经预训练的模型
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
# 在CoCo上加载经过预训练的预训练模型
model=torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# 将分类器替换为具有用户定义的num_classes的新分类器
num_classes=2#1 (person)+background
# 获取分类器的输入参数的数量
in_features=model.roi_heads.box_predictor.cls_score.in_features
# 用新的头部替换预先训练好的头部
model.roi_heads.box_predictor=FastRCNNPredictor(in_features, num_classes)
#%% 修改模型以添加不同的主干
import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
# 加载预先训练的模型进行分类和返回
backbone=torchvision.models.mobilenet_v2(pretrained=True).features
# faster r-cnn需要知道骨干网络中的输出通道数量，对于mobilenet_v2,是1280，这里添加backbone.out_channels=1280
backbone.out_channels=1280
# 让RPN在每个空间位置生成5*3个锚点，具有5种不同的大小和3种不同的宽高比
# 有一个元组，每个特征映射可能具有不同的大小和宽高比
anchor_generator=AnchorGenerator(sizes=((32,64,128,256,512),),
aspect_ratios=((0.5,1.0,2.0),))
# 定义用于执行区域裁剪的特征映射，以及重新缩放后裁剪的大小。如果主干返回Tensor,featmap_names应为[0]更一般的，主干应该返回OrderedDict[Tensor],并在feature_names中，选择要使用的功能映射
roi_pooler=torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
# 将这些pieces放在FasterRCNN模型中
model=FastRCNNPredictor(backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
我们从预先训练的模型进行微调，因为数据集非常小，所以将遵循上述第一种情况。还要计算实例分割掩膜，因此将使用Mask R-CNN:
#%%定义模型
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
def get_model_instance_segmentation(num_classes):
# 加载在CoCo上预训练的实例分割模型
model=torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
# 获取分类器的输入参数的数量
in_features=model.roi_heads.box_predictor.cls_score.in_features
# 用新的头部替换预先训练好的头部
model.roi_heads.box_predictor=FastRCNNPredictor(in_features, num_classes)
# 获取掩膜分类器的输入特征数
in_features_mask=model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer=256
# 用新的掩膜预测期替换掩膜预测期
model.roi_heads.mask_predictor=MaskRCNNPredictor(in_features_mask,
hidden_layer,
num_classes)
return model
在**reference/detection/**中有许多辅助函数来简化训练和评估检测模型。此处将使用reference/detection/engine.py，reference/detection/util.py和reference/detection/transforms.py。只需要将他们复制到文件夹并在工作路径使用——>我在这里
#为数据扩充/转换编写辅助函数
import sys
sys.path.append("D:\Python\Pytorch")
import transforms as t
def get_transform(train):
transforms=[]
transforms.append(t.ToTensor())
if train:
transforms.append(t.RandomHorizontalFlip(0.5))
return t.Compose(transforms)
#编写执行训练和验证的主要功能
from engine import train_one_epoch,evaluate
import utils
def main():
# 在GPU训练，没有使用CPU
device=torch.device('cuda')if torch.cuda.is_available() else torch.device('cpu')
# 数据集只有两类人：背景和人
num_classes=2
# 使用数据集和定义的转换
dataset=PennFudanDataset('PennFudanPed', get_transform(train=True))
dataset_test=PennFudanDataset('PennFudanPed', get_transform(train=False))
# 在训练和测试集拆分数据集
indices=torch.randperm(len(dataset)).tolist()
dataset=torch.utils.data.Subset(dataset,indices[:-50])
dataset_test=torch.utils.data.Subset(dataset_test, indices[-50:])
# 定义训练和验证数据加载器
data_loader=torch.utils.data.DataLoader(
dataset,batch_size=2,shuffle=True,num_workers=4,
collate_fn=utils.collate_fn)
data_loader_test=torch.utils.data.DataLoader(
dataset,batch_size=1,shuffle=False,num_workers=4,
collate_fn=utils.collate_fn)
# 使用辅助函数获取模型
model=get_model_instance_segmentation(num_classes)
# 将模型迁移到合适的设备
model.to(device)
# 构造一个优化器
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)
# 和学习率调度程序
lr_scheduler=torch.optim.lr_scheduler.StepLR(optimizer, step_size=3,gamma=0.1)
# 训练10个epoch
num_epochs=10
for epoch in range(num_epochs):
# 训练一个epoch，每10此迭代打印一次
train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
# 更新学习速率
lr_scheduler.step()
# 在测试集上评价
evaluate(model, data_loader_test, device=device)
print("That‘s all")
在第一个epoch训练后可以得到一下结果：
因此，在一个epoch训练之后，获得了CoCo-style mAP为64.8，mask mAP为67.6——如下图
经过训练10个epoch后，得到如下指标：
小结明天再写，太晚了，我要赶紧回宿舍洗洗睡了。明天见，猿友们~