fix cpu count

training/coqui_stt_training/evaluate.py

@@ -4,7 +4,7 @@
 
 import json
 import sys
-from multiprocessing import cpu_count
+import psutil
 
 import progressbar
 import tensorflow.compat.v1 as tfv1
@@ -91,7 +91,7 @@ def evaluate(test_csvs, create_model):
 
     # Get number of accessible CPU cores for this process
     try:
-        num_processes = cpu_count()
+        num_processes = len(psutil.Process().cpu_affinity())
     except NotImplementedError:
         num_processes = 1
 

training/coqui_stt_training/evaluate_export.py

@@ -8,7 +8,8 @@
 import wave
 import io
 from functools import partial
-from multiprocessing import JoinableQueue, Manager, Process, cpu_count
+from multiprocessing import JoinableQueue, Manager, Process
+import psutil
 
 import numpy as np
 from coqui_stt_training.util.evaluate_tools import calculate_and_print_report
@@ -142,7 +143,7 @@ def parse_args():
     parser.add_argument(
         "--proc",
         required=False,
-        default=cpu_count(),
+        default=len(psutil.Process().cpu_affinity()),
         type=int,
         help="Number of processes to spawn, defaulting to number of CPUs",
     )

training/coqui_stt_training/evaluate_flashlight.py

@@ -4,7 +4,7 @@
 
 import json
 import sys
-from multiprocessing import cpu_count
+import psutil
 
 import progressbar
 import tensorflow.compat.v1 as tfv1
@@ -95,7 +95,7 @@ def evaluate(test_csvs, create_model):
 
     # Get number of accessible CPU cores for this process
     try:
-        num_processes = cpu_count()
+        num_processes = len(psutil.Process().cpu_affinity())
     except NotImplementedError:
         num_processes = 1
 

training/coqui_stt_training/evaluate_wav2vec2am.py

@@ -8,8 +8,9 @@
 import os
 import sys
 from functools import partial
-from multiprocessing import JoinableQueue, Manager, Process, cpu_count
+from multiprocessing import JoinableQueue, Manager, Process
 from pathlib import Path
+import psutil
 
 import numpy as np
 import onnxruntime
@@ -274,7 +275,7 @@ def parse_args():
     parser.add_argument(
         "--proc",
         required=False,
-        default=cpu_count(),
+        default=len(psutil.Process().cpu_affinity()),
         type=int,
         help="Number of processes to spawn, defaulting to number of CPUs",
     )

training/coqui_stt_training/transcribe.py

@@ -9,6 +9,7 @@
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Optional, List, Tuple
+import psutil
 
 LOG_LEVEL_INDEX = sys.argv.index("--log_level") + 1 if "--log_level" in sys.argv else 0
 DESIRED_LOG_LEVEL = (
@@ -34,7 +35,7 @@
 
 
 def cpu_count():
-    return os.cpu_count() or 1
+    return len(psutil.Process().cpu_affinity()) or 1
 
 
 class TranscriptionPool(PoolBase):

training/coqui_stt_training/util/augmentations.py

@@ -1,8 +1,8 @@
 import math
-import os
 import random
 import re
 from multiprocessing import Process, Queue
+import psutil
 
 import numpy as np
 import resampy
@@ -312,7 +312,14 @@ def __repr__(self):
 
     def start(self, buffering=BUFFER_SIZE):
         self.queue = Queue(
-            max(1, math.floor(self.probability * self.layers[1] * os.cpu_count()))
+            max(
+                1,
+                math.floor(
+                    self.probability
+                    * self.layers[1]
+                    * len(psutil.Process().cpu_affinity())
+                ),
+            )
         )
         self.enqueue_process = Process(
             target=_enqueue_overlay_samples,

training/coqui_stt_training/util/helpers.py

@@ -1,5 +1,6 @@
 import heapq
 import os
+import psutil
 import random
 import sys
 import time
@@ -134,7 +135,14 @@ def __init__(
         process_ahead=None,
         sleeping_for=0.1,
     ):
-        self.process_ahead = os.cpu_count() if process_ahead is None else process_ahead
+        if processes is None:
+            processes = len(psutil.Process().cpu_affinity())
+
+        self.process_ahead = (
+            len(psutil.Process().cpu_affinity())
+            if process_ahead is None
+            else process_ahead
+        )
         self.sleeping_for = sleeping_for
         self.processed = 0
         self.pool = Pool(

training/coqui_stt_training/util/lm_optimize_wav2vec2am.py

@@ -2,9 +2,9 @@
 # -*- coding: utf-8 -*-
 from __future__ import absolute_import, print_function
 
-import os
 import sys
 from dataclasses import dataclass, field
+import psutil
 
 import optuna
 from clearml import Task
@@ -103,7 +103,7 @@ class LmOptimizeWav2vec2amConfig(BaseSttConfig):
         ),
     )
     num_processes: int = field(
-        default=os.cpu_count(),
+        default=len(psutil.Process().cpu_affinity()),
         metadata=dict(help="Number of worker processes for evaluation."),
     )
     clearml_project: str = field(

training/coqui_stt_training/util/multiprocessing.py

@@ -4,6 +4,7 @@
 import os
 import sys
 from contextlib import contextmanager
+import psutil
 
 
 def target_fn(*args, **kwargs):
@@ -39,7 +40,7 @@ class PoolBase:
     all child processes in order to synchronize work between all processes. You
     can also use `self.process_id`, which is an integer, unique per process,
     increasing in value from 0 to processes-1 (if not specified, processes
-    defaults to os.cpu_count()).
+    defaults to len(psutil.Process().cpu_affinity()) ).
 
     `run` will be called, in the child processes, potentially multiple times, in
     order to process data.
@@ -70,7 +71,7 @@ def run(self, x):
     @classmethod
     def create_impl(cls, processes=None, context=None, initargs=(), *args, **kwargs):
         if processes is None:
-            processes = os.cpu_count()
+            processes = len(psutil.Process().cpu_affinity())
 
         if context is None:
             context = multiprocessing