多卡预处理

bert_gen.py

@@ -3,7 +3,7 @@ from multiprocessing import Pool
 import commons
 import utils
 from tqdm import tqdm
-
+from loguru import logger
 from text import cleaned_text_to_sequence, get_bert
 
 config_path = "configs/config.json"
@@ -12,6 +12,12 @@ total_gpus = torch.cuda.device_count()
 
 
 def process_line(line):
+    rank = mp.current_process()._identity
+    rank = rank[0] if len(rank) > 0 else 0
+    if torch.cuda.is_available():
+        gpu_id = rank % torch.cuda.device_count()
+        device = torch.device(f"cuda:{gpu_id}")
+    logger.info(f"Rank {rank} uses device {device}")
     wav_path, _, language_str, text, phones, tone, word2ph = line.strip().split("|")
     phone = phones.split(" ")
     tone = [int(i) for i in tone.split(" ")]
@@ -33,7 +39,7 @@ def process_line(line):
         bert = torch.load(bert_path)
         assert bert.shape[-1] == len(phone)
     except Exception:
-        bert = get_bert(text, word2ph, language_str)
+        bert = get_bert(text, word2ph, language_str,device)
         assert bert.shape[-1] == len(phone)
         torch.save(bert, bert_path)