### 如何使用 Detectron2 训练 COCO 数据集 #### 准备环境和安装依赖库 为了能够顺利运行 Detectron2 并训练 COCO 数据集，需先设置好 Python 环境并安装必要的包。确保已安装 PyTorch 和 torchvision 版本兼容于所使用的 CUDA 版本。 ```bash pip install git+https://github.com/facebookresearch/detectron2.git ``` #### 下载与预处理 COCO 数据集 按照官方指导，可以获取 COCO 数据集，并将其放置在一个合适的位置以便后续加载。对于 COCO 数据集而言，通常会有一个图像文件夹以及对应的标注 JSON 文件[^2]。 ```python import os from detectron2.data.datasets import register_coco_instances from detectron2.config import get_cfg from detectron2.engine import DefaultTrainer, default_argument_parser, launch # 注册数据集到 Detectron2 的 DatasetCatalog 中 register_coco_instances("my_dataset_train", {}, "path/to/json/train.json", "path/to/image_dir") register_coco_instances("my_dataset_val", {}, "path/to/json/val.json", "path/to/image_dir") # 验证注册是否成功 print("Dataset registered.") ``` #### 设置配置参数 (Config) 创建或修改配置文件来指定模型架构、超参数和其他选项。这里展示了一个简单的例子，其中设置了 `num_gpus` 参数用于控制 GPU 数量，并指定了要使用的预训练权重路径[^1]。 ```python cfg = get_cfg() cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")) cfg.DATASETS.TRAIN = ("my_dataset_train",) cfg.DATASETS.TEST = () # no metrics implemented for this dataset cfg.DATALOADER.NUM_WORKERS = 2 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml") # Let training initialize from model zoo cfg.SOLVER.IMS_PER_BATCH = 2 cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR cfg.SOLVER.MAX_ITER = 300 # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset (default: 512) cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class (ballon). (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets) os.makedirs(cfg.OUTPUT_DIR, exist_ok=True) ``` #### 开始训练过程 一旦完成了上述准备工作，则可以通过调用默认 Trainer 类来进行实际的训练工作。下面给出了一段完整的脚本来启动多GPU分布式训练任务。 ```python if __name__ == "__main__": args = default_argument_parser().parse_args() trainer = DefaultTrainer(cfg) trainer.resume_or_load(resume=False) launch( trainer.train, num_gpus_per_machine=cfg.num_gpus, dist_url="tcp://localhost:8089", ) ```