42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
from typing import List
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
|
|
class SequencePacker:
|
|
"""序列打包策略"""
|
|
|
|
def __init__(self, pack_size: int, pad_value: int = 0):
|
|
self.pack_size = pack_size
|
|
self.pad_value = pad_value
|
|
|
|
def pack(self, sequences: List[Tensor]) -> List[Tensor]:
|
|
"""打包序列到固定大小"""
|
|
packages = []
|
|
sequences.sort(key=lambda x: x.numel(), reverse=True)
|
|
|
|
current_pack = torch.full((self.pack_size,), self.pad_value, dtype=torch.int32)
|
|
current_pos = 0
|
|
|
|
for tensor in sequences:
|
|
tensor = tensor[:self.pack_size] if tensor.numel() > self.pack_size else tensor
|
|
tensor_size = tensor.numel()
|
|
|
|
if current_pos + tensor_size > self.pack_size:
|
|
packages.append(current_pack)
|
|
current_pack = torch.full((self.pack_size,), self.pad_value, dtype=torch.int32)
|
|
current_pos = 0
|
|
|
|
current_pack[current_pos:current_pos + tensor_size] = tensor
|
|
current_pos += tensor_size
|
|
|
|
if current_pos > 0:
|
|
packages.append(current_pack)
|
|
|
|
return packages
|
|
|
|
|
|
def pack_sequences(sequences: List[Tensor], pack_size: int, pad_value: int) -> List[Tensor]:
|
|
"""向后兼容的函数接口"""
|
|
return SequencePacker(pack_size, pad_value).pack(sequences)
|