YOLOX指定GPU训练

1. 找到YOLOX-main/yolox/utils/dist.py文件

注释以下代码：

def get_local_rank() -> int:
    """
    Returns:
        The rank of the current process within the local (per-machine) process group.
    """
    if _LOCAL_PROCESS_GROUP is None:
        return get_rank()

    if not dist.is_available():
        return 0
    if not dist.is_initialized():
        return 0
    return dist.get_rank(group=_LOCAL_PROCESS_GROUP)

添加以下代码：

def get_local_rank() -> int:              ##########add
    """
    Returns:
        The rank of the current process within the local (per-machine) process group.
    """
    if not dist.is_available():
        return select_device('1')  //设置哪块gpu就修改此处，比如select_device('0')、select_device('1')
    if not dist.is_initialized():
        return select_device('1')  //设置哪块gpu就修改此处，比如select_device('0')、select_device('1')
    assert _LOCAL_PROCESS_GROUP is not None
    return dist.get_rank(group=_LOCAL_PROCESS_GROUP)


def select_device(device=''):
    # device = 'cpu' or '0' or '0,1,2,3'
    cpu_request = device.lower() == 'cpu'
    if device and not cpu_request:  # if device requested other than 'cpu'
        os.environ['CUDA_VISIBLE_DEVICES'] = device  # set environment variable
        assert torch.cuda.is_available(), 'CUDA unavailable, invalid device %s requested' % device  # check availablity

    cuda = False if cpu_request else torch.cuda.is_available()

    return 0 if cuda else 'cpu'         ##########add

2. 找到train.py文件

修改以下代码中的default，原为0，现改为1

    parser.add_argument(
        "-d", "--devices", default=1, type=int, help="device for training"  
    )

3. 修改训练指令

原为 python3 tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 0 -b 32 -c yolox_s.pth
现为 python3 tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 1 -b 32 -c yolox_s.pth

资料来源于网络
参考链接 https://blog.csdn.net/weixin_44245653/article/details/124804793