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

这段代码主要是用于分布式训练的，首先会获取当前设备中 CUDA 设备的数量，然后如果需要进行分布式训练，就会使用 `dist.init_process_group` 函数初始化分布式训练环境，使用 `os.environ` 获取当前进程的 `LOCAL_RANK` 和 `RANK`，并将设备设置为当前进程的 CUDA 设备。如果不需要进行分布式训练，则将设备设置为 CUDA 设备（如果可用），否则设置为 CPU 设备，并将 `local_rank` 和 `rank` 设置为 0。最后，如果当前进程是主进程（即 `local_rank` 为 0），则输出当前进程的信息和 CUDA 设备的数量。