From d02113267b69222896c0978338387e13dd6204a5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= Date: Tue, 5 Sep 2023 19:34:28 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A4=9A=E5=8D=A1=E9=A2=84=E5=A4=84=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bert_gen.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) 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)