ngpus_per_node = torch.cuda.device_count() if distributed: dist.init_process_group(backend="nccl") local_rank = int(os.environ["LOCAL_RANK"]) rank = int(os.environ["RANK"]) device = torch.device("cuda", local_rank) if local_rank == 0: print(f"[{os.getpid()}] (rank = {rank}, local_rank = {local_rank}) training...") print("Gpu Device Count : ", ngpus_per_node) else: device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') local_rank = 0 rank = 0
时间: 2023-12-31 14:06:49 浏览: 170
torch.cuda.is_available()返回False解决方案
5星 · 资源好评率100%
这段代码看起来是用于设置分布式训练的GPU设备使用情况。首先通过 `torch.cuda.device_count()` 获取当前节点上可用的GPU数量,然后通过 `dist.init_process_group()` 初始化分布式环境,其中 `backend` 参数指定使用的通信后端,这里是使用了 `nccl`。接下来通过获取环境变量中的 `LOCAL_RANK` 和 `RANK` 参数,确定当前进程的本地rank和全局rank。如果本地rank为0,则输出日志信息,并打印可用的GPU数量。如果不是分布式训练,则直接判断CUDA是否可用,选择使用CPU或GPU。
阅读全文