mmrotate在cpu上运行的方法

最新推荐文章于 2023-07-13 00:16:51 发布

shenhaibb_

最新推荐文章于 2023-07-13 00:16:51 发布

阅读量913

点赞数

分类专栏： Python 文章标签：深度学习 python 机器学习

本文链接：https://blog.csdn.net/shenhaibb_/article/details/127145924

版权

Python 专栏收录该内容

14 篇文章 2 订阅

订阅专栏

修改如下代码：

mmrotate/apis/train.py Line 60

# put model on gpus
    if distributed:
        find_unused_parameters = cfg.get('find_unused_parameters', False)
        # Sets the `find_unused_parameters` parameter in
        # torch.nn.parallel.DistributedDataParallel
        model = MMDistributedDataParallel(
            model.cuda(),
            device_ids=[torch.cuda.current_device()],
            broadcast_buffers=False,
            find_unused_parameters=find_unused_parameters)
    else:
        # model = MMDataParallel(model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)
        model = MMDataParallel(model.cpu(), device_ids=cfg.gpu_ids)

最后两行 model.cuda()换成model.cpu()