reafactor: 重构项目
This commit is contained in:
+94
-73
@@ -1,109 +1,130 @@
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
from typing import List
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
from .utils import error_handler
|
||||
from pipeline.utils import error_handler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SequencePacker:
|
||||
"""
|
||||
Packs variable-length sequences into fixed-size tensors, suitable for
|
||||
concatenating unequal-length training samples into uniform shapes
|
||||
for DataLoader / model training.
|
||||
|
||||
def __init__(self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = torch.int32):
|
||||
Algorithm (Sorted Greedy Fill, based on First-Fit Decreasing heuristic):
|
||||
|
||||
Input: sequences = [A(len=5), B(len=2), C(len=3)], pack_size = 8
|
||||
|
||||
1. Validate & Normalize
|
||||
- Check 1D dimension, unify dtype, truncate overlong sequences with warning
|
||||
- Result: [(A,5), (B,2), (C,3)]
|
||||
|
||||
2. Sort by length descending (FFD)
|
||||
- Result: [(A,5), (C,3), (B,2)]
|
||||
|
||||
3. Greedy fill: write into a pre-allocated buffer sequentially, flush when full
|
||||
- Write A(5) -> buffer = [A A A A A _ _ _], pos=5
|
||||
- Write C(3) -> pos+3=8 <= 8 -> buffer = [A A A A A C C C], pos=8
|
||||
- Buffer full -> flush as package[0], reset buffer & pos=0
|
||||
- Write B(2) -> buffer = [B B _ _ _ _ _ _], pos=2
|
||||
- Loop ends -> flush tail -> package[1] = [B B 0 0 0 0 0 0]
|
||||
|
||||
Output: [package[0], package[1]]
|
||||
|
||||
Cross-group consistency:
|
||||
When packing different key groups (e.g. sequences and loss_masks)
|
||||
with separate pack() calls, tensors at the same index always have
|
||||
identical lengths, so the descending sort produces the exact same
|
||||
ordering. Element-level correspondence across groups is preserved.
|
||||
|
||||
Performance:
|
||||
- Pre-allocated buffer reused via fill_() to avoid repeated tensor creation
|
||||
- Attributes cached as local variables inside the loop to reduce lookup overhead
|
||||
"""
|
||||
|
||||
def __init__(self, pack_size: int, pad_value: int = 0, dtype: torch.dtype = None):
|
||||
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."""
|
||||
# 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.dtype = dtype # None = follow input dtype
|
||||
self._buffer: Tensor | None = None
|
||||
self._pos = 0
|
||||
self._packages: List[Tensor] = []
|
||||
# Backward compatibility: maintain _current_pack reference
|
||||
self._current_pack = self._buffer
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset packer state for instance reuse, unlocking dtype."""
|
||||
self.dtype = None
|
||||
self._buffer = None
|
||||
self._pos = 0
|
||||
self._packages = []
|
||||
|
||||
@error_handler()
|
||||
def pack(self, sequences: List[Tensor]) -> List[Tensor]:
|
||||
"""
|
||||
Pack sequences into fixed-size packages.
|
||||
Pack sequences into fixed-size packages using First-Fit Decreasing.
|
||||
|
||||
Sequences are sorted by length descending to minimize wasted padding.
|
||||
All tensor groups (e.g. sequences, loss_masks) with matching per-item
|
||||
lengths produce identical ordering, so cross-group correspondence is preserved.
|
||||
|
||||
Args:
|
||||
sequences: List of input tensors
|
||||
sequences: List of 1D input tensors.
|
||||
|
||||
Returns:
|
||||
List of packed tensors, each with length equal to pack_size
|
||||
List of packed tensors, each with length equal to pack_size.
|
||||
"""
|
||||
# Input validation
|
||||
if not sequences:
|
||||
return []
|
||||
|
||||
# Validate and cache tensor sizes in one pass
|
||||
tensor_sizes = []
|
||||
|
||||
# --- validate & normalize in a single pass ---
|
||||
normalized: list[tuple[Tensor, int]] = []
|
||||
target_dtype = self.dtype if self.dtype is not None else sequences[0].dtype
|
||||
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}"
|
||||
)
|
||||
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."
|
||||
)
|
||||
if seq.dtype != target_dtype:
|
||||
seq = seq.to(target_dtype)
|
||||
length = seq.numel()
|
||||
if length > self.pack_size:
|
||||
seq = seq[: self.pack_size]
|
||||
length = self.pack_size
|
||||
normalized.append((seq, length))
|
||||
|
||||
# Reset state for new packing
|
||||
# --- reset internal state ---
|
||||
buf = self._buffer
|
||||
if buf is None or buf.dtype != target_dtype:
|
||||
buf = torch.full((self.pack_size,), self.pad_value, dtype=target_dtype)
|
||||
self._buffer = buf
|
||||
buf.fill_(self.pad_value)
|
||||
self._pos = 0
|
||||
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, tensor_size in indexed_seqs:
|
||||
# Truncate sequences that exceed pack_size
|
||||
if tensor_size > self.pack_size:
|
||||
logger.warning(
|
||||
f"Sequence length {tensor_size} exceeds pack_size {self.pack_size}, truncating"
|
||||
)
|
||||
tensor_size = self.pack_size
|
||||
tensor = tensor[: self.pack_size]
|
||||
# --- sort by length descending (FFD heuristic) ---
|
||||
normalized.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
# Current package is full, create a new one
|
||||
if self._current_pos + tensor_size > self.pack_size:
|
||||
# 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
|
||||
# --- greedy fill ---
|
||||
buf = self._buffer
|
||||
pos = self._pos
|
||||
packages = self._packages
|
||||
pack_size = self.pack_size
|
||||
pad_value = self.pad_value
|
||||
|
||||
# Place tensor in current package
|
||||
self._buffer[self._current_pos : self._current_pos + tensor_size] = tensor
|
||||
self._current_pos += tensor_size
|
||||
for tensor, length in normalized:
|
||||
if pos + length > pack_size:
|
||||
# flush current package
|
||||
packages.append(buf.clone())
|
||||
buf.fill_(pad_value)
|
||||
pos = 0
|
||||
buf[pos : pos + length] = tensor
|
||||
pos += length
|
||||
|
||||
# Handle the last package (pad to pack_size)
|
||||
if self._current_pos > 0:
|
||||
package = self._buffer.clone()
|
||||
self._packages.append(package)
|
||||
# flush the last (possibly partial) package
|
||||
if pos > 0:
|
||||
packages.append(buf.clone())
|
||||
|
||||
# Clear buffer and reset state for backward compatibility
|
||||
self._buffer = None
|
||||
self._current_pack = None
|
||||
self._current_pos = 0
|
||||
|
||||
# write back state
|
||||
self._pos = pos
|
||||
return self._packages
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset packer state for reuse. More efficient than creating a new instance."""
|
||||
self._reset()
|
||||
Reference in New Issue
Block a user