diff --git a/bert_gen.py b/bert_gen.py index f71ba77..8616800 100644 --- a/bert_gen.py +++ b/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)