Detectron2项目启动：COCO数据集操作、代码编写及环境疑问

一、cudatoolkit版本选择与安装 #

你的CUDA版本是11.8，直接安装cudatoolkit 11.8即可，版本完全匹配才能让PyTorch和CUDA正常联动：

• 用conda安装：conda install cudatoolkit=11.8 -c nvidia
• 用pip安装：pip install nvidia-cudatoolkit==11.8.0

二、完成Detectron2环境部署 #

1. 进入你下载的Detectron2代码根目录
2. 执行本地安装命令：pip install -e .
   这条命令会自动安装Detectron2依赖的所有库（比如pycocotools、opencv-python等）

三、自定义COCO数据集的目录结构要求 #

确保你的数据集按以下结构存放：

your_dataset/
├── annotations/
│   ├── train.json  # 训练集标注文件
│   └── val.json    # 验证集标注文件
├── train/          # 训练集图片文件夹
│   ├── img1.jpg
│   └── ...
└── val/            # 验证集图片文件夹
    ├── img2.jpg
    └── ...

四、核心.py文件内容示例（以训练为例） #

以下是一个完整的训练脚本框架，包含数据集注册、配置加载、训练器初始化等关键部分：

import os
from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg
from detectron2.data.datasets import register_coco_instances
from detectron2.data import MetadataCatalog, DatasetCatalog

# 1. 注册自定义COCO数据集
register_coco_instances("your_dataset_train", {}, 
                        "/path/to/your_dataset/annotations/train.json", 
                        "/path/to/your_dataset/train")
register_coco_instances("your_dataset_val", {}, 
                        "/path/to/your_dataset/annotations/val.json", 
                        "/path/to/your_dataset/val")

# 2. 获取数据集元数据（可选，用于可视化）
train_metadata = MetadataCatalog.get("your_dataset_train")
train_dataset_dicts = DatasetCatalog.get("your_dataset_train")

# 3. 配置模型与训练参数
cfg = get_cfg()
# 选择预训练模型配置，根据任务类型替换（比如分割任务选COCO-InstanceSegmentation下的配置）
cfg.merge_from_file("detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")  
cfg.DATASETS.TRAIN = ("your_dataset_train",)
cfg.DATASETS.TEST = ("your_dataset_val",)
cfg.DATALOADER.NUM_WORKERS = 4
# 加载预训练权重，路径可替换为本地预训练文件
cfg.MODEL.WEIGHTS = "detectron2://COCO-Detection/faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl"  
cfg.SOLVER.IMS_PER_BATCH = 2  # 根据GPU显存调整，显存不足时调小
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 30000  # 根据数据集大小调整，小数据集可设为10000-20000
cfg.SOLVER.STEPS = []        # 不需要学习率衰减则留空
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 5  # 替换为你的数据集类别数
cfg.OUTPUT_DIR = "./output"  # 训练结果保存目录

# 4. 初始化训练器并启动训练
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

五、关键操作说明与参考内容 #

1. 配置文件选择：Detectron2的configs目录下有各类任务（检测、分割等）的预训练配置，根据需求选择对应子目录下的文件。
2. 参数调整：
   • SOLVER.IMS_PER_BATCH：单GPU显存不足时调小（比如1），多GPU训练时按比例增加
   • MODEL.ROI_HEADS.NUM_CLASSES：必须和你的数据集类别数完全一致
   • MAX_ITER：小数据集可适当降低迭代次数，避免过拟合
3. 推理脚本示例：如果需要做推理，替换训练部分代码为以下逻辑：

   from detectron2.engine import DefaultPredictor
   from detectron2.utils.visualizer import Visualizer
   import cv2
   
   cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
   cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # 推理置信度阈值
   predictor = DefaultPredictor(cfg)
   im = cv2.imread("/path/to/test/image.jpg")
   outputs = predictor(im)
   # 可视化推理结果
   v = Visualizer(im[:, :, ::-1], metadata=train_metadata, scale=1.0)
   out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
   cv2.imwrite("result.jpg", out.get_image()[:, :, ::-1])
   

4. 核心参考文档内容：
   • Quick Start：涵盖环境搭建、基本训练推理流程
   • Custom Datasets章节：详细说明自定义COCO数据集的注册方法
   • Configs章节：解释各配置参数的含义与调整规则

内容的提问来源于stack exchange，提问作者Evian Jo