添加多卡预处理

Author: colstone

basics/base_binarizer.py

@@ -310,6 +310,23 @@ def process_dataset(self, prefix, num_workers=0, apply_augmentation=False):
 
         aug_map = self.arrange_data_augmentation(self.meta_data_iterator(prefix)) if apply_augmentation else {}
 
+        device_ids = None
+        num_workers = int(num_workers)
+        if (
+                num_workers > 0 and torch.cuda.is_available()
+                and (torch.cuda.device_count() > 1)
+        ):
+            per_gpu_workers = self.binarization_args.get('num_workers_per_gpu')
+            if per_gpu_workers is None and self.binarization_args.get('workers_per_gpu', False):
+                per_gpu_workers = num_workers
+            if per_gpu_workers:
+                per_gpu_workers = int(per_gpu_workers)
+                device_ids = [
+                    gpu for gpu in range(torch.cuda.device_count())
+                    for _ in range(per_gpu_workers)
+                ]
+                num_workers = len(device_ids)
+
         def postprocess(_item):
             nonlocal total_sec, total_raw_sec, extra_info, max_no
             if _item is None:
@@ -349,7 +366,9 @@ def postprocess(_item):
             if num_workers > 0:
                 # code for parallel processing
                 for item in tqdm(
-                        chunked_multiprocess_run(self.process_item, args, num_workers=num_workers),
+                        chunked_multiprocess_run(
+                            self.process_item, args, num_workers=num_workers, device_ids=device_ids
+                        ),
                         total=len(list(self.meta_data_iterator(prefix)))
                 ):
                     postprocess(item)

configs/original/base.yaml

@@ -11,6 +11,7 @@ binarizer_cls: null
 binarization_args:
   shuffle: false
   num_workers: 0
+  workers_per_gpu: false
 
 audio_sample_rate: 44100
 hop_size: 512

modules/pe/__init__.py

@@ -5,13 +5,13 @@
 from .rmvpe import RMVPE
 
 
-def initialize_pe():
+def initialize_pe(device=None):
     pe = hparams['pe']
     pe_ckpt = hparams['pe_ckpt']
     if pe == 'parselmouth':
         return ParselmouthPE()
     elif pe == 'rmvpe':
-        return RMVPE(pe_ckpt)
+        return RMVPE(pe_ckpt, device=device)
     elif pe == 'harvest':
         return HarvestPE()
     else:

modules/pe/rmvpe/inference.py

@@ -13,9 +13,11 @@
 
 
 class RMVPE(BasePE):
-    def __init__(self, model_path, hop_length=160):
+    def __init__(self, model_path, hop_length=160, device=None):
         self.resample_kernel = {}
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = torch.device(device)
         self.model = E2E0(4, 1, (2, 2)).eval().to(self.device)
         ckpt = torch.load(model_path, map_location=self.device)
         self.model.load_state_dict(ckpt['model'], strict=False)

preprocessing/acoustic_binarizer.py

@@ -134,7 +134,7 @@ def process_item(self, item_name, meta_data, binarization_args):
         # get ground truth f0
         global pitch_extractor
         if pitch_extractor is None:
-            pitch_extractor = initialize_pe()
+            pitch_extractor = initialize_pe(self.device)
         gt_f0, uv = pitch_extractor.get_pitch(
             waveform, samplerate=hparams['audio_sample_rate'], length=length,
             hop_size=hparams['hop_size'], f0_min=hparams['f0_min'], f0_max=hparams['f0_max'],
@@ -229,7 +229,7 @@ def arrange_data_augmentation(self, data_iterator):
         aug_list = []
         all_item_names = [item_name for item_name, _ in data_iterator]
         total_scale = 0
-        aug_pe = initialize_pe()
+        aug_pe = initialize_pe(self.device)
         if self.augmentation_args['random_pitch_shifting']['enabled']:
             from augmentation.spec_stretch import SpectrogramStretchAugmentation
             aug_args = self.augmentation_args['random_pitch_shifting']

preprocessing/variance_binarizer.py

@@ -300,7 +300,7 @@ def process_item(self, item_name, meta_data, binarization_args):
 
         global pitch_extractor
         if pitch_extractor is None:
-            pitch_extractor = initialize_pe()
+            pitch_extractor = initialize_pe(self.device)
         f0 = uv = None
         if self.prefer_ds:
             f0_seq = self.load_attr_from_ds(ds_id, name, 'f0_seq', idx=ds_seg_idx)

utils/multiprocess_utils.py

@@ -12,7 +12,15 @@ def main_process_print(self, *args, sep=' ', end='\n', file=None):
         print(self, *args, sep=sep, end=end, file=file)
 
 
-def chunked_worker_run(map_func, args, results_queue=None):
+def chunked_worker_run(map_func, args, results_queue=None, device_id=None):
+    if device_id is not None:
+        try:
+            import torch
+            torch.cuda.set_device(device_id)
+            if hasattr(map_func, '__self__') and map_func.__self__ is not None:
+                map_func.__self__.device = torch.device(f'cuda:{device_id}')
+        except Exception:
+            traceback.print_exc()
     for a in args:
         # noinspection PyBroadException
         try:
@@ -25,10 +33,15 @@ def chunked_worker_run(map_func, args, results_queue=None):
             results_queue.put(None)
 
 
-def chunked_multiprocess_run(map_func, args, num_workers, q_max_size=1000):
+def chunked_multiprocess_run(map_func, args, num_workers, q_max_size=1000, device_ids=None):
     num_jobs = len(args)
     if num_jobs < num_workers:
         num_workers = num_jobs
+        if device_ids is not None:
+            device_ids = device_ids[:num_workers]
+
+    if device_ids is not None:
+        assert len(device_ids) == num_workers
 
     queues = [Manager().Queue(maxsize=q_max_size // num_workers) for _ in range(num_workers)]
     if platform.system().lower() != 'windows':
@@ -39,7 +52,9 @@ def chunked_multiprocess_run(map_func, args, num_workers, q_max_size=1000):
     workers = []
     for i in range(num_workers):
         worker = process_creation_func(
-            target=chunked_worker_run, args=(map_func, args[i::num_workers], queues[i]), daemon=True
+            target=chunked_worker_run,
+            args=(map_func, args[i::num_workers], queues[i], None if device_ids is None else device_ids[i]),
+            daemon=True
         )
         workers.append(worker)
         worker.start()