### 使用 PyTorch 实现 Faster R-CNN 训练自定义数据集 为了使用 PyTorch 的 `FasterRCNN` 模型训练自定义数据集，需遵循一系列特定步骤来准备环境、加载并预处理数据以及配置模型参数。 #### 准备工作 安装必要的库和依赖项： ```bash pip install torch torchvision torchaudio ``` 确保已克隆或下载了适用于 PyTorch 的 Faster R-CNN 代码仓库[^2]。 #### 加载与预处理数据 创建一个继承自 `torch.utils.data.Dataset` 类的数据集类用于读取图像及其对应的标注文件。假设这些标签是以 COCO JSON 文件的形式给出，则可以利用 `pycocotools` 库解析它们。 ```python from pycocotools.coco import COCO import os from PIL import Image class CustomDataset(torchvision.datasets.CocoDetection): def __init__(self, root, annotation, transforms=None): super().__init__(root=root, annFile=annotation) self.transforms = transforms def __getitem__(self, idx): img_id = self.ids[idx] path = self.coco.loadImgs(img_id)[0]['file_name'] image = Image.open(os.path.join(self.root, path)).convert('RGB') annotations_ids = self.coco.getAnnIds(imgIds=img_id) annotations = self.coco.loadAnns(annotations_ids) boxes = [] labels = [] area = [] for i in range(len(annotations)): bbox = annotations[i]['bbox'] label = annotations[i]['category_id'] xmin = bbox[0] ymin = bbox[1] xmax = xmin + bbox[2] ymax = ymin + bbox[3] boxes.append([xmin, ymin, xmax, ymax]) labels.append(label) area.append(bbox[2]*bbox[3]) target = {} target["boxes"] = torch.as_tensor(boxes, dtype=torch.float32) target["labels"] = torch.as_tensor(labels, dtype=torch.int64) target["image_id"] = torch.tensor([img_id]) target["area"] = torch.as_tensor(area, dtype=torch.float32) if self.transforms is not None: image, target = self.transforms(image, target) return image, target def __len__(self): return len(self.ids) ``` 此部分实现了自定义数据集的初始化方法及获取单个样本的方法，在其中完成了图片路径定位、边界框坐标提取等工作，并返回转换后的张量形式的结果以便后续网络输入。 #### 配置模型 接下来构建基于 ResNet50-FPN 的 Faster R-CNN 检测器实例化过程如下所示: ```python def get_model_instance_segmentation(num_classes): model = torchvision.models.detection.fasterrcnn_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) return model ``` 这里先调用了带有预训练权重的基础版本(`pretrained=True`)，再修改分类头以适应新类别数量的需求[^1]。 #### 开始训练 最后编写简单的训练循环脚本执行批量迭代更新操作直至收敛为止；期间还需注意保存最佳性能下的检查点(checkpoint)，方便之后恢复继续优化或是部署应用阶段直接加载使用。 ```python device = 'cuda' if torch.cuda.is_available() else 'cpu' 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) num_epochs = 10 for epoch in range(num_epochs): train_one_epoch(model, optimizer, data_loader_train, device, epoch, print_freq=10) evaluate(model, data_loader_val, device=device) ``` 上述代码片段展示了完整的训练流程设置，包括设备选择、优化算法指定、损失函数计算等环节[^4]。