Pytorch 入门学习(检测)
数据集从标准的类torch.utils.data.Dataset继承,并实现_len_, getitem。
唯一要求的特性是数据集__getitem__应该返回:
图像:PIL图像大小(H,W)
目标:包含以下字段的字典
<1> boxes(FloatTensor[N,4]):N边框(bounding boxes)坐标的格式[x0,x1,y0,y1],取值范围是0到W,0到H。
<2> labels(Int64Tensor[N]):每个边框的标签。
<3> image_id(Int64Tensor1):图像识别器,它应该在数据集中的所有图像中是唯一的,并在评估期间使用。
<4> area(Tensor[N]):边框的面积,在使用COCO指标进行评估时使用此项来分隔小、中和大框之间的度量标准得分。
<5> iscrowed(UInt8Tensor[N,H,W]):在评估期间属性设置为iscrowed=True的实例会被忽略。
<6> (可选)masks(UInt8Tesor[N,H,W]):每个对象的分段掩码。
<7> (可选)keypoints (FloatTensor[N, K, 3]:对于N个对象中的每一个,它包含[x,y,visibility]格式的K个关键点,用 于定义对象。visibility = 0表示关键点不可见。请注意,对于数据扩充,翻转关键点的概念取决于数据表示,您应该调整 reference/detection/transforms.py 以用于新的关键点表示。
如果你的模型返回上述方法,它们将使其适用于培训和评估,并将使用 pycocotools 的评估脚本。
此外,如果要在训练期间使用宽高比分组(以便每个批次仅包含具有相似宽高比的图像),则建议还实现get_height_and_width方法, 该方法返回图像的高度和宽度。如果未提供此方法,我们将通过__getitem__查询数据集的所有元素,这会将图像加载到内存中,但比提供自定义方法时要慢。
有两种常见情况可能需要修改torchvision modelzoo中的一个可用模型。第一个是我们想要从预先训练的模型开始,然后微调最后一层。 另一种是当我们想要用不同的模型替换模型的主干时(例如,用于更快的预测)。
1 微调已经预训练的模型 让我们假设你想从一个在COCO上已预先训练过的模型开始,并希望为你的特定类进行微调。这是一种可行的方法:
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
# 在COCO上加载经过预训练的预训练模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# replace the classifier with a new one, that has
# 将分类器替换为具有用户定义的 num_classes的新分类器
num_classes = 2 # 1 class (person) + background
# 获取分类器的输入参数的数量
in_features = model.roi_heads.box_predictor.cls_score.in_features
# 用新的头部替换预先训练好的头部
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
2 修改模型以添加不同的主干
import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
# 加载预先训练的模型进行分类和返回
# 只有功能
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
# FasterRCNN需要知道骨干网中的输出通道数量。对于mobilenet_v2,它是1280,所以我们需要在这里添加它
backbone.out_channels = 1280
# 我们让RPN在每个空间位置生成5 x 3个锚点
# 具有5种不同的大小和3种不同的宽高比。
# 我们有一个元组[元组[int]]
# 因为每个特征映射可能具有不同的大小和宽高比
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
# 定义一下我们将用于执行感兴趣区域裁剪的特征映射,以及重新缩放后裁剪的大小。
# 如果您的主干返回Tensor,则featmap_names应为[0]。
# 更一般地,主干应该返回OrderedDict [Tensor]
# 并且在featmap_names中,您可以选择要使用的功能映射。
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
# 将这些pieces放在FasterRCNN模型中
model = FasterRCNN(backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
在references/detection/中,我们有许多辅助函数来简化训练和评估检测模型。在这里,我们将使用 references/detection/engine.py,references/detection/utils.py和references/detection/transforms.py。 只需将它们复制到您的文件夹并在此处使用它们。
from engine import train_one_epoch, evaluate
import utils
def main():
# 在GPU上训练,若无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_test, 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个epochs
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 it!")
有关包含multi-machine / multi-gpu training的更完整示例,请检查 torchvision 存储库中的references/detection/train.py。
参考链接:
https://github.com/fendouai/PyTorchDocs/blob/master/fourSection/ObjectDetectionFinetuning.md