PyTorch学习笔记(十二) ---- PyTorch目标检测

悠悠 2023-02-14 10:15 1阅读 0赞

**转载请注明作者和出处：** [http://blog.csdn.net/john\_bh/][http_blog.csdn.net_john_bh]

### 文章目录 ###

*  1.定义数据集
 *  2.为 PennFudan 编写自定义数据集
 *  3. 定义模型
 *   *  3.1 PennFudan 数据集的实例分割模型
 *  4.训练和验证 完整代码
 *   *  4.1 为数据扩充/转换编写辅助函数：
     *  4.2 编写执行训练和验证的主要功能
     *  4.3. 完整代码
 *  5.总结

在文中将微调在 Penn-Fudan 数据库中对行人检测和分割的已预先训练的 Mask R-CNN 模型。它包含170个图像和345个行人实例，用它来说明如何在 torchvision 中使用新功能，以便在自定义数据集上训练实例分割模型。

另外，在 [references / detection][references _ detection] 中，有许多辅助函数来简化训练和评估检测模型。这次学习中，需要使用 [references/detection/engine.py][references_detection_engine.py] ，[references/detection/utils.py][references_detection_utils.py] 和 [references/detection/transforms.py][references_detection_transforms.py]， 需将它们复制到文件夹下。

# 1.定义数据集 #

对于训练对象检测的引用脚本，实例分割和人员关键点检测要求能够轻松支持添加新的自定义数据。数据集应该从标准的类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(Int64Tensor\[1\])：图像识别器，它应该在数据集中的所有图像中是唯一的，并在评估期间使用。
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\_\_查询数据集的所有元素，这会将图像加载到内存中，**但比提供自定义方法时要慢。**

# 2.为 PennFudan 编写自定义数据集 #

[PennFudan数据集下载连接][PennFudan]  
文件结构：

> PennFudanPed/  
> PedMasks/  
> FudanPed00001\_mask.png  
> FudanPed00002\_mask.png  
> FudanPed00003\_mask.png  
> FudanPed00004\_mask.png  
> …  
> PNGImages/  
> FudanPed00001.png  
> FudanPed00002.png  
> FudanPed00003.png  
> FudanPed00004.png

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 1]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 2]  
因此每个图像具有相应的分割掩膜，其中每个颜色对应于不同的实例。接下来为这个数据集写一个torch.utils.data.Dataset类

为数据扩充/转换编写辅助函数：

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)

# 3. 定义模型 #

现在定义一个可以上述数据集执行预测的模型。使用 [Mask R-CNN][]， 它基于 [Faster R-CNN][]。Faster R-CNN 是一种模型，可以预测图像中潜在对象的边界框和类别得分。  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 3]  
Mask R-CNN 在 Faster R-CNN 中添加了一个额外的分支，它还预测每个实例的分割蒙版。  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 4]  
有两种常见情况可能需要修改 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)

1.  修改模型以添加不同的主干

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)

## 3.1 PennFudan 数据集的实例分割模型 ##

因为数据集非常小，所以将遵循上述第一种情况：从预先训练的模型中进行微调。另外还要计算实例分割掩膜，将使用 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

这将使模型准备好的自定义数据集上进行训练和评估。

# 4.训练和验证 完整代码 #

在`references/detection/`中，有许多辅助函数来简化训练和评估检测模型。在这里，将使用 `references/detection/engine.py`，`references/detection/utils.py`和`references/detection/transforms.py`。 只需将它们复制到自己的文件夹并在此处使用它们。

## 4.1 为数据扩充/转换编写辅助函数： ##

## 4.2 编写执行训练和验证的主要功能 ##

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!")

在第一个epoch训练后可以得到下面的结果：

> Epoch: \[0\] \[ 0/60\] eta: 0:01:18 lr: 0.000090 loss: 2.5213 (2.5213) loss\_classifier: 0.8025 (0.8025) loss\_box\_reg: 0.2634 (0.2634) loss\_mask: 1.4265 (1.4265) loss\_objectness: 0.0190 (0.0190) loss\_rpn\_box\_reg: 0.0099 (0.0099) time: 1.3121 data: 0.3024 max mem: 3485  
> Epoch: \[0\] \[10/60\] eta: 0:00:20 lr: 0.000936 loss: 1.3007 (1.5313) loss\_classifier: 0.3979 (0.4719) loss\_box\_reg: 0.2454 (0.2272) loss\_mask: 0.6089 (0.7953) loss\_objectness: 0.0197 (0.0228) loss\_rpn\_box\_reg: 0.0121 (0.0141) time: 0.4198 data: 0.0298 max mem: 5081  
> Epoch: \[0\] \[20/60\] eta: 0:00:15 lr: 0.001783 loss: 0.7567 (1.1056) loss\_classifier: 0.2221 (0.3319) loss\_box\_reg: 0.2002 (0.2106) loss\_mask: 0.2904 (0.5332) loss\_objectness: 0.0146 (0.0176) loss\_rpn\_box\_reg: 0.0094 (0.0123) time: 0.3293 data: 0.0035 max mem: 5081  
> Epoch: \[0\] \[30/60\] eta: 0:00:11 lr: 0.002629 loss: 0.4705 (0.8935) loss\_classifier: 0.0991 (0.2517) loss\_box\_reg: 0.1578 (0.1957) loss\_mask: 0.1970 (0.4204) loss\_objectness: 0.0061 (0.0140) loss\_rpn\_box\_reg: 0.0075 (0.0118) time: 0.3403 data: 0.0044 max mem: 5081  
> Epoch: \[0\] \[40/60\] eta: 0:00:07 lr: 0.003476 loss: 0.3901 (0.7568) loss\_classifier: 0.0648 (0.2022) loss\_box\_reg: 0.1207 (0.1736) loss\_mask: 0.1705 (0.3585) loss\_objectness: 0.0018 (0.0113) loss\_rpn\_box\_reg: 0.0075 (0.0112) time: 0.3407 data: 0.0044 max mem: 5081  
> Epoch: \[0\] \[50/60\] eta: 0:00:03 lr: 0.004323 loss: 0.3237 (0.6703) loss\_classifier: 0.0474 (0.1731) loss\_box\_reg: 0.1109 (0.1561) loss\_mask: 0.1658 (0.3201) loss\_objectness: 0.0015 (0.0093) loss\_rpn\_box\_reg: 0.0093 (0.0116) time: 0.3379 data: 0.0043 max mem: 5081  
> Epoch: \[0\] \[59/60\] eta: 0:00:00 lr: 0.005000 loss: 0.2540 (0.6082) loss\_classifier: 0.0309 (0.1526) loss\_box\_reg: 0.0463 (0.1405) loss\_mask: 0.1568 (0.2945) loss\_objectness: 0.0012 (0.0083) loss\_rpn\_box\_reg: 0.0093 (0.0123) time: 0.3489 data: 0.0042 max mem: 5081  
> Epoch: \[0\] Total time: 0:00:21 (0.3570 s / it)  
> creating index…  
> index created!  
> Test: \[ 0/50\] eta: 0:00:19 model\_time: 0.2152 (0.2152) evaluator\_time: 0.0133 (0.0133) time: 0.4000 data: 0.1701 max mem: 5081  
> Test: \[49/50\] eta: 0:00:00 model\_time: 0.0628 (0.0687) evaluator\_time: 0.0039 (0.0064) time: 0.0735 data: 0.0022 max mem: 5081  
> Test: Total time: 0:00:04 (0.0828 s / it)  
> Averaged stats: model\_time: 0.0628 (0.0687) evaluator\_time: 0.0039 (0.0064)  
> Accumulating evaluation results…  
> DONE (t=0.01s).  
> Accumulating evaluation results…  
> DONE (t=0.01s).  
> IoU metric: bbox  
> Average Precision (AP) @\[ IoU=0.50:0.95 | area= all | maxDets=100 \] = 0.606  
> Average Precision (AP) @\[ IoU=0.50 | area= all | maxDets=100 \] = 0.984  
> Average Precision (AP) @\[ IoU=0.75 | area= all | maxDets=100 \] = 0.780  
> Average Precision (AP) @\[ IoU=0.50:0.95 | area= small | maxDets=100 \] = 0.313  
> Average Precision (AP) @\[ IoU=0.50:0.95 | area=medium | maxDets=100 \] = 0.582  
> Average Precision (AP) @\[ IoU=0.50:0.95 | area= large | maxDets=100 \] = 0.612  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area= all | maxDets= 1 \] = 0.270  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area= all | maxDets= 10 \] = 0.672  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area= all | maxDets=100 \] = 0.672  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area= small | maxDets=100 \] = 0.650  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area=medium | maxDets=100 \] = 0.755  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area= large | maxDets=100 \] = 0.664  
> IoU metric: segm  
> Average Precision (AP) @\[ IoU=0.50:0.95 | area= all | maxDets=100 \] = 0.704  
> Average Precision (AP) @\[ IoU=0.50 | area= all | maxDets=100 \] = 0.979  
> Average Precision (AP) @\[ IoU=0.75 | area= all | maxDets=100 \] = 0.871  
> Average Precision (AP) @\[ IoU=0.50:0.95 | area= small | maxDets=100 \] = 0.325  
> Average Precision (AP) @\[ IoU=0.50:0.95 | area=medium | maxDets=100 \] = 0.488  
> Average Precision (AP) @\[ IoU=0.50:0.95 | area= large | maxDets=100 \] = 0.727  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area= all | maxDets= 1 \] = 0.316  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area= all | maxDets= 10 \] = 0.748  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area= all | maxDets=100 \] = 0.749  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area= small | maxDets=100 \] = 0.650  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area=medium | maxDets=100 \] = 0.673  
> Average Recall (AR) @\[ IoU=0.50:0.95 | area= large | maxDets=100 \] = 0.758

因此，在一个epoch训练之后，我们获得了COCO-style mAP为60.6，并且mask mAP为70.4。经过训练10个epoch后，我得到了以下指标：

预测结果:  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 5]  
训练的模型预测了此图像中的9个人物，让我们看看其中的几个，由下图可以看到预测效果很好。  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 6]

## 4.3. 完整代码 ##

# Sample code from the TorchVision 0.3 Object Detection Finetuning Tutorial
    # http://pytorch.org/tutorials/intermediate/torchvision_tutorial.html
    
    import os
    import numpy as np
    import torch
    from PIL import Image
    
    import torchvision
    from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
    from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
    
    from engine import train_one_epoch, evaluate
    import utils
    import transforms as T
    
    
    class PennFudanDataset(object):
        def __init__(self, root, transforms):
            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 ad 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)
    
            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])
    
            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])
            area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
            # suppose all instances are not crowd
            iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
    
            target = {
        }
            target["boxes"] = boxes
            target["labels"] = labels
            target["masks"] = masks
            target["image_id"] = image_id
            target["area"] = area
            target["iscrowd"] = iscrowd
    
            if self.transforms is not None:
                img, target = self.transforms(img, target)
    
            return img, target
    
        def __len__(self):
            return len(self.imgs)
    
    def get_model_instance_segmentation(num_classes):
        # load an instance segmentation model pre-trained pre-trained on COCO
        model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
    
        # 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)
    
        # now get the number of input features for the mask classifier
        in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
        hidden_layer = 256
        # and replace the mask predictor with a new one
        model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
                                                           hidden_layer,
                                                           num_classes)
    
        return model
    
    
    def get_transform(train):
        transforms = []
        transforms.append(T.ToTensor())
        if train:
            transforms.append(T.RandomHorizontalFlip(0.5))
        return T.Compose(transforms)
    
    
    def main():
        # train on the GPU or on the CPU, if a GPU is not available
        device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    
        # our dataset has two classes only - background and person
        num_classes = 2
        # use our dataset and defined transformations
        dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
        dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))
    
        # split the dataset in train and test set
        indices = torch.randperm(len(dataset)).tolist()
        dataset = torch.utils.data.Subset(dataset, indices[:-50])
        dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
    
        # define training and validation data loaders
        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)
    
        # get the model using our helper function
        model = get_model_instance_segmentation(num_classes)
    
        # move model to the right device
        model.to(device)
    
        # 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)
    
        # 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)
    
        print("That's it!")
        
    if __name__ == "__main__":
        main()

# 5.总结 #

在本教程中，您学习了如何在自定义数据集上为实例分段模型创建自己的训练管道。为此，您编写了一个torch.utils.data.Dataset类， 它返回图像以及地面实况框和分割掩码。您还利用了在COCO train2017上预训练的Mask R-CNN模型，以便对此新数据集执行传输学习。

有关包含`multi-machine / multi-gpu training`的更完整示例，请检查 torchvision 存储库中的`references/detection/train.py`。

[http_blog.csdn.net_john_bh]: http://blog.csdn.net/john_bh/
[references _ detection]: https://github.com/pytorch/vision/tree/master/references/detection
[references_detection_engine.py]: https://github.com/pytorch/vision/blob/master/references/detection/engine.py
[references_detection_utils.py]: https://github.com/pytorch/vision/blob/master/references/detection/utils.py
[references_detection_transforms.py]: https://github.com/pytorch/vision/blob/master/references/detection/transforms.py
[PennFudan]: https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center]: https://img-blog.csdnimg.cn/20200602102750214.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg=,size_16,color_FFFFFF,t_70#pic_center
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 1]: https://img-blog.csdnimg.cn/20200602102703476.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg=,size_16,color_FFFFFF,t_70#pic_center
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 2]: https://img-blog.csdnimg.cn/20200602102712445.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg=,size_16,color_FFFFFF,t_70#pic_center
[Mask R-CNN]: https://arxiv.org/abs/1703.06870
[Faster R-CNN]: https://arxiv.org/abs/1506.01497
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 3]: https://img-blog.csdnimg.cn/20200602105854300.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg=,size_16,color_FFFFFF,t_70#pic_center
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 4]: https://img-blog.csdnimg.cn/20200602105914575.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg=,size_16,color_FFFFFF,t_70#pic_center
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 5]: https://img-blog.csdnimg.cn/20200822160805747.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg=,size_16,color_FFFFFF,t_70#pic_center
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg_size_16_color_FFFFFF_t_70_pic_center 6]: https://img-blog.csdnimg.cn/20200822160827172.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2pvaG5fYmg=,size_16,color_FFFFFF,t_70#pic_center