在多个 GPU 上运行detectron2 训练

Question

我在多个 GPU 上运行修改后的 train_net.py 脚本时遇到问题。

重现问题的说明：

我正在使用this dataset 作为实验来测试如何使用 Slurm 在多个 GPU 上运行 detectron2 训练。

1. 完整的可运行代码或您所做的全部更改（tools/train_net.py 已修改）：

#!/usr/bin/env python
# Copyright (c) Facebook, Inc. and its affiliates.
"""
A main training script.

This scripts reads a given config file and runs the training or evaluation.
It is an entry point that is made to train standard models in detectron2.

In order to let one script support training of many models,
this script contains logic that are specific to these built-in models and therefore
may not be suitable for your own project.
For example, your research project perhaps only needs a single "evaluator".

Therefore, we recommend you to use detectron2 as an library and take
this file as an example of how to use the library.
You may want to write your own script with your datasets and other customizations.
"""

import logging
import os
from collections import OrderedDict
import torch

import detectron2.utils.comm as comm
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import get_cfg
from detectron2.data import MetadataCatalog
from detectron2.engine import DefaultPredictor, DefaultTrainer, default_argument_parser, default_setup, hooks, launch
from detectron2.evaluation import (
    CityscapesInstanceEvaluator,
    CityscapesSemSegEvaluator,
    COCOEvaluator,
    COCOPanopticEvaluator,
    DatasetEvaluators,
    LVISEvaluator,
    PascalVOCDetectionEvaluator,
    SemSegEvaluator,
    verify_results,
)
from detectron2.modeling import GeneralizedRCNNWithTTA

from detectron2 import model_zoo
from detectron2.data.datasets import register_coco_instances
from detectron2.utils.visualizer import Visualizer
import glob
import cv2

# Dodato za SLURM
import ssl
ssl._create_default_https_context = ssl._create_unverified_context
cfg = ""
trainer = ""

class Trainer(DefaultTrainer):
    """
    We use the "DefaultTrainer" which contains pre-defined default logic for
    standard training workflow. They may not work for you, especially if you
    are working on a new research project. In that case you can write your
    own training loop. You can use "tools/plain_train_net.py" as an example.
    """
    @classmethod

    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
        if output_folder is None:
            os.makedirs("coco_eval", exist_ok=True)
            output_folder = "coco_eval"
        return COCOEvaluator(dataset_name, cfg, False, output_folder)

    @classmethod
    def test_with_TTA(cls, cfg, model):
        logger = logging.getLogger("detectron2.trainer")
        # In the end of training, run an evaluation with TTA
        # Only support some R-CNN models.
        logger.info("Running inference with test-time augmentation ...")
        model = GeneralizedRCNNWithTTA(cfg, model)
        evaluators = [
            cls.build_evaluator(
                cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
            )
            for name in cfg.DATASETS.TEST
        ]
        res = cls.test(cfg, model, evaluators)
        res = OrderedDict({k + "_TTA": v for k, v in res.items()})
        return res


def setup():
    global cfg
    """
    Create configs and perform basic setups.
    """
    print("START SETUP")
    cfg = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
    cfg.DATASETS.TRAIN = ("my_dataset_train",)
    cfg.DATASETS.TEST = ("my_dataset_val",)
    cfg.DATALOADER.NUM_WORKERS = 4
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml")  # Let training initialize from model zoo
    cfg.SOLVER.IMS_PER_BATCH = 2
    cfg.SOLVER.BASE_LR = 0.001
    cfg.SOLVER.WARMUP_ITERS = 50
    cfg.SOLVER.MAX_ITER = 500 #adjust up if val mAP is still rising, adjust down if overfit
    cfg.SOLVER.STEPS = (50, 450)
    cfg.SOLVER.GAMMA = 0.05
    cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 32
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4 #your number of classes + 1
    cfg.TEST.EVAL_PERIOD = 500
    # cfg.merge_from_list(args.opts)
    # cfg.freeze()
    default_setup(cfg, args)
    print("END SETUP")
    return cfg


def main():
    global cfg, trainer
    cfg = setup()

    if args.eval_only:
        print("EVAL_ONLY")
        model = Trainer.build_model(cfg)
        DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(cfg.MODEL.WEIGHTS, resume=args.resume)
        res = Trainer.test(cfg, model)
        if cfg.TEST.AUG.ENABLED:
            res.update(Trainer.test_with_TTA(cfg, model))
        if comm.is_main_process():
            verify_results(cfg, res)
        return res
    print("BEFORE TRAINER")
    trainer = Trainer(cfg)
    trainer.resume_or_load(resume=args.resume)
    if cfg.TEST.AUG.ENABLED:
        print("TEST AUG ENABLED")
        trainer.register_hooks([hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))])
    print("BEFORE MAIN END")
    return trainer.train()


if __name__ == "__main__":
    from datetime import datetime

    args = default_argument_parser().parse_args()
    print("Command Line Args:", args)
    register_coco_instances("my_dataset_train", {}, "./data/train/_annotations.coco.json", "./data/train")
    register_coco_instances("my_dataset_val", {}, "./data/valid/_annotations.coco.json", "./data/valid")
    register_coco_instances("my_dataset_test", {}, "./data/test/_annotations.coco.json", "./data/test")

    now = datetime.now()
    print("before launch =", now)

    launch(main, num_gpus_per_machine = args.num_gpus, dist_url = "auto")

    now = datetime.now()
    print("after launch =", now)

    # Evaluate

    from detectron2.data import DatasetCatalog, build_detection_test_loader
    from detectron2.evaluation import inference_on_dataset

    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.85
    predictor = DefaultPredictor(cfg)
    evaluator = COCOEvaluator("my_dataset_test", cfg, False, output_dir="./output/")
    val_loader = build_detection_test_loader(cfg, "my_dataset_test")
    inference_on_dataset(trainer.model, val_loader, evaluator)

    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
    cfg.DATASETS.TEST = ("my_dataset_test", )
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set the testing threshold for this model
    predictor = DefaultPredictor(cfg)
    test_metadata = MetadataCatalog.get("my_dataset_test")

    with open("results.txt", mode="a") as f:
        for imageName in glob.glob('data/test/*jpg'):
            im = cv2.imread(imageName)
            outputs = predictor(im)
            f.write(f"Instances:{outputs['instances']}\n")
            v = Visualizer(im[:, :, ::-1], metadata=test_metadata)
            out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
            cv2.imwrite(f"images/{imageName}",out.get_image()[:, :, ::-1])

另一方面，我有这个 slurm 脚本可以在 2 个 GPU 上运行实验：

#!/bin/bash -l

#SBATCH --account=Account
#SBATCH --partition=gpu # gpu partition
#SBATCH --nodes=1 # 1 node, 4 GPUs per node
#SBATCH --time=24:00:00 
#SBATCH --job-name=detectron2_demo4 # job name

module load Python/3.9.5-GCCcore-10.3.0
module load CUDA/11.1.1-GCC-10.2.0

cd /experiment_path

srun python main.py --num-gpus 2

当我运行这个脚本时，我遇到了一个错误（cat slurm-xxx.out），并且没有错误文件：

The following have been reloaded with a version change:
  1) GCCcore/10.3.0 => GCCcore/10.2.0
  2) binutils/2.36.1-GCCcore-10.3.0 => binutils/2.35-GCCcore-10.2.0
  3) zlib/1.2.11-GCCcore-10.3.0 => zlib/1.2.11-GCCcore-10.2.0

Command Line Args: Namespace(config_file='', resume=False, eval_only=False, num_gpus=2, num_machines=1, machine_rank=0, dist_url='tcp://127.0.0.1:54190', opts=[])
before launch = 2021-08-03 21:26:48.061817

[W ProcessGroupNCCL.cpp:1569] Rank 0 using best-guess GPU 0 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect.Specify device_ids in barrier() to force use of a particular device.

[W ProcessGroupNCCL.cpp:1569] Rank 1 using best-guess GPU 1 to perform barrier as devices used by this process are currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect.Specify device_ids in barrier() to force use of a particular device.

预期行为：

在 2 个 GPU 上运行训练

环境：

粘贴以下命令的输出：

No CUDA runtime is found, using CUDA_HOME='/usr/local/software/CUDAcore/11.1.1'
---------------------  --------------------------------------------------------------------------------
sys.platform           linux
Python                 3.9.5 (default, Jul  9 2021, 09:35:24) [GCC 10.3.0]
numpy                  1.21.1
detectron2             0.5 @/home/users/aimhigh/detectron2/detectron2
Compiler               GCC 10.2
CUDA compiler          CUDA 11.1
DETECTRON2_ENV_MODULE  <not set>
PyTorch                1.9.0+cu102 @/home/users/aimhigh/.local/lib/python3.9/site-packages/torch
PyTorch debug build    False
GPU available          No: torch.cuda.is_available() == False
Pillow                 8.3.1
torchvision            0.10.0+cu102 @/home/users/aimhigh/.local/lib/python3.9/site-packages/torchvision
fvcore                 0.1.5.post20210727
iopath                 0.1.9
cv2                    4.5.3
---------------------  --------------------------------------------------------------------------------
PyTorch built with:
  - GCC 7.3
  - C++ Version: 201402
  - Intel(R) Math Kernel Library Version 2020.0.0 Product Build 20191122 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v2.1.2 
  - OpenMP 201511 (a.k.a. OpenMP 4.5)
  - NNPACK is enabled
  - CPU capability usage: AVX2

提示（更多细节）：

Githubdetectron2：link

Github pyTorch：link

Github NCCL：link

Answer 1

这是 Slurm 文件的问题。 所以，我通过在 Slurm 文件中附加此代码来解决它：

export NCCL_SOCKET_IFNAME=bond0
export NCCL_IB_DISABLE=1
export NCCL_P2P_DISABLE=1