Pytorch入门

原文链接: https://github.com/fendouai/PyTorchDocs/blob/master/fourSection/ObjectDetectionFinetuning.md

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

你可能感兴趣的:(Pytorch入门)