perf: 优化 processors、cache、packing 模块性能并简化 README

pipeline/packing.py

@@ -1,5 +1,5 @@
 import logging
-from typing import List
+from typing import List, Optional
 import torch
 from torch import Tensor
 
@@ -14,14 +14,23 @@ class SequencePacker:
         self.pack_size = pack_size
         self.pad_value = pad_value
         self.dtype = dtype
+        # Pre-allocate buffer for better performance
+        self._buffer: Optional[Tensor] = None
         self._reset()
 
     def _reset(self) -> None:
         """Reset internal state for instance reuse."""
-        self._current_pack = torch.full(
-            (self.pack_size,), self.pad_value, dtype=self.dtype
-        )
+        # Reuse buffer instead of creating new tensors
+        if self._buffer is None or self._buffer.shape[0] != self.pack_size:
+            self._buffer = torch.full(
+                (self.pack_size,), self.pad_value, dtype=self.dtype
+            )
+        else:
+            self._buffer.fill_(self.pad_value)
         self._current_pos = 0
+        self._packages: List[Tensor] = []
+        # Backward compatibility: maintain _current_pack reference
+        self._current_pack = self._buffer
 
     @error_handler()
     def pack(self, sequences: List[Tensor]) -> List[Tensor]:
@@ -37,53 +46,63 @@ class SequencePacker:
         # Input validation
         if not sequences:
             return []
+        
+        # Validate and cache tensor sizes in one pass
+        tensor_sizes = []
         for i, seq in enumerate(sequences):
             if seq.dim() != 1:
                 raise ValueError(
                     f"Expected 1D tensor at index {i}, got {seq.dim()}D tensor with shape {seq.shape}"
                 )
-            # Check dtype compatibility and warn if mismatched
+            tensor_sizes.append(seq.numel())
             if seq.dtype != self.dtype:
                 logger.warning(
                     f"Input tensor dtype {seq.dtype} does not match packer dtype {self.dtype}, "
                     f"will be converted. This may affect packing efficiency."
                 )
 
-        packages = []
-        # Sort by length in descending order to improve packing efficiency
-        # Use sorted() to avoid modifying the input list
-        sorted_sequences = sorted(sequences, key=lambda x: x.numel(), reverse=True)
+        # Reset state for new packing
+        self._packages = []
+        self._reset()
+        
+        # Combine sequences with their sizes for sorting
+        indexed_seqs = list(zip(sequences, tensor_sizes))
+        # Sort by size descending (First-Fit Decreasing algorithm)
+        indexed_seqs.sort(key=lambda x: x[1], reverse=True)
 
-        for tensor in sorted_sequences:
+        for tensor, tensor_size in indexed_seqs:
             # Truncate sequences that exceed pack_size
-            if tensor.numel() > self.pack_size:
+            if tensor_size > self.pack_size:
                 logger.warning(
-                    f"Sequence length {tensor.numel()} exceeds pack_size {self.pack_size}, truncating"
+                    f"Sequence length {tensor_size} exceeds pack_size {self.pack_size}, truncating"
                 )
+                tensor_size = self.pack_size
                 tensor = tensor[: self.pack_size]
-            tensor_size = tensor.numel()
 
             # Current package is full, create a new one
             if self._current_pos + tensor_size > self.pack_size:
-                packages.append(self._current_pack)
-                self._current_pack = torch.full(
-                    (self.pack_size,), self.pad_value, dtype=self.dtype
-                )
+                # Finish current package (pad to pack_size)
+                package = self._buffer.clone()
+                self._packages.append(package)
+                # Reset buffer for reuse
+                self._buffer.fill_(self.pad_value)
                 self._current_pos = 0
 
-            # Place tensor in current package (remaining positions stay as pad_value)
-            self._current_pack[self._current_pos : self._current_pos + tensor_size] = (
-                tensor
-            )
+            # Place tensor in current package
+            self._buffer[self._current_pos : self._current_pos + tensor_size] = tensor
             self._current_pos += tensor_size
 
-        # Handle the last package
+        # Handle the last package (pad to pack_size)
         if self._current_pos > 0:
-            packages.append(self._current_pack)
-            self._current_pack = None
-            self._current_pos = 0
+            package = self._buffer.clone()
+            self._packages.append(package)
 
-        return packages
+        # Clear buffer and reset state for backward compatibility
+        self._buffer = None
+        self._current_pack = None
+        self._current_pos = 0
+        
+        return self._packages
 
     def reset(self) -> None:
         """Reset packer state for reuse. More efficient than creating a new instance."""