tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Update train_ms.py. Auto choose gloo backend while using Windows platform. (#134)

Browse Source
This commit is contained in:
Sora
2023-10-30 20:31:22 +08:00
committed by GitHub
parent b13c0fad9d
commit 88efbed72e
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
train_ms.py
View File

@@ -1,5 +1,6 @@
# flake8: noqa: E402 # flake8: noqa: E402
import platform
import os import os
import torch import torch
from torch.nn import functional as F from torch.nn import functional as F
@@ -53,8 +54,11 @@ def run():
os.environ[env_name] = str(env_value) os.environ[env_name] = str(env_value)
# 多卡训练设置 # 多卡训练设置
backend = "nccl"
if platform.system() == "Windows":
backend = "gloo"
dist.init_process_group( dist.init_process_group(
backend="nccl", backend=backend,
init_method="env://", # If Windows,switch to gloo backend. init_method="env://", # If Windows,switch to gloo backend.
) # Use torchrun instead of mp.spawn ) # Use torchrun instead of mp.spawn
rank = dist.get_rank() rank = dist.get_rank()