Upload folder using huggingface_hub

online_bspline_tokenizer.py

@@ -24,6 +24,36 @@ from functools import wraps
 from transformers.processing_utils import ProcessorMixin
 from scipy.interpolate import BSpline, make_lsq_spline
 
+# ============================================================================
+# Process Pool Manager (Singleton)
+# ============================================================================
+class ProcessPoolManager:
+    """Singleton process pool manager to avoid repeated pool creation/destruction."""
+    _instance = None
+    _pool = None
+    _lock = mp.Lock()
+    
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+    
+    def get_pool(self, processes=None):
+        """Get or create the process pool."""
+        with self._lock:
+            if self._pool is None:
+                processes = processes or mp.cpu_count()
+                self._pool = mp.Pool(processes=processes)
+        return self._pool
+    
+    def close_pool(self):
+        """Close and cleanup the process pool."""
+        with self._lock:
+            if self._pool is not None:
+                self._pool.close()
+                self._pool.join()
+                self._pool = None
+
 # ============================================================================
 # Utility Functions
 # ============================================================================
@@ -216,13 +246,21 @@ class BestBSpline:
         joint_dof: Number of joint DOFs
         gripper_dof: Number of gripper DOFs (从后往前数)
         check_step: Downsampling step for constraint checking acceleration
+        is_multi_process: Whether to use multiprocessing for batch compression (default: False)
     """
-    def __init__(self, degree: int = 3, joint_dof: int = 6, gripper_dof: int = 1, check_step=1, use_gurobi: bool = False):
+    def __init__(self, degree: int = 3, joint_dof: int = 6, gripper_dof: int = 1, check_step=1, use_gurobi: bool = False, is_multi_process: bool = False):
         self.degree = degree
         self.joint_dof = joint_dof
         self.gripper_dof = gripper_dof
         self.check_step = check_step   # 降采样步长，用于加速
         self.use_gurobi = use_gurobi
+        self.is_multi_process = is_multi_process
+        self.pool_manager = ProcessPoolManager() if is_multi_process else None
+
+    def __del__(self):
+        """Cleanup process pool when object is destroyed."""
+        if self.is_multi_process and self.pool_manager is not None:
+            self.pool_manager.close_pool()
 
     def _extract_gripper_forced_knots(self, gripper_traj: np.ndarray, time_points: np.ndarray) -> List[int]:
         """
@@ -400,15 +438,18 @@ class BestBSpline:
     def batch_compress(self, batch_trajectory: torch.Tensor):
         batch = batch_trajectory.shape[0]
         cpu_cores_per_fit = max(1, mp.cpu_count() // batch)
-        # with mp.Pool(processes=batch) as pool:
-        #     tasks = [(batch_trajectory[i], cpu_cores_per_fit) for i in range(batch)]
-        #     results = pool.starmap(self.fit, tasks)
 
-        # NOTE：单线程测试是否会出现问题
-        results = []
-        for i in batch_trajectory:
-            res = self.fit(i, cpu_cores=cpu_cores_per_fit)
-            results.append(res)
+        if self.is_multi_process and batch > 1:
+            # Use singleton process pool for parallel processing
+            pool = self.pool_manager.get_pool(processes=batch)
+            tasks = [(batch_trajectory[i], cpu_cores_per_fit) for i in range(batch)]
+            results = pool.starmap(self.fit, tasks)
+        else:
+            # Single-threaded processing
+            results = []
+            for i in batch_trajectory:
+                res = self.fit(i, cpu_cores=cpu_cores_per_fit)
+                results.append(res)
 
         batch_knots = [res[0] for res in results]
         batch_control_points = [res[1] for res in results]
@@ -477,6 +518,7 @@ class BestTokenizer(torch.nn.Module, ProcessorMixin):
         degree: B-spline degree (default 3 = cubic)
         gripper_zero_order: Use separate zero-order spline for gripper (True/False)
         gripper_dof: Number of gripper DOFs (if gripper_zero_order is True)
+        is_multi_process: Whether to use multiprocessing for batch compression (default: False)
     
     Example:
         >>> tokenizer = BestTokenizer(num_dof=7, num_basis=10, seq_len=50)
@@ -499,7 +541,7 @@ class BestTokenizer(torch.nn.Module, ProcessorMixin):
 
     def __init__(self, num_dof: int = 7, in_seq_len: int = 10, out_seq_len: int = 5,
                  vocab_size: int = 256, degree: int = 3, gripper_dof: int = 1, 
-                 do_pad: bool = True, use_gurobi: bool = False, device: str = "cuda"):
+                 do_pad: bool = True, use_gurobi: bool = False, is_multi_process: bool = False, device: str = "cuda"):
         super().__init__()
         self.in_seq_len = in_seq_len
         self.out_seq_len = out_seq_len 
@@ -516,6 +558,7 @@ class BestTokenizer(torch.nn.Module, ProcessorMixin):
             joint_dof=self.joint_dof,
             gripper_dof=self.gripper_dof,
             use_gurobi=use_gurobi,
+            is_multi_process=is_multi_process,
         )
         
         # Initialize weight bounds for normalization

processor_config.json

@@ -11,5 +11,6 @@
   "gripper_dof": 1,
   "do_pad": true,
   "use_gurobi": false,
+  "is_multi_process": false,
   "device": "cuda"
 }