36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
from typing import List
|
||
import torch
|
||
from torch import Tensor
|
||
|
||
|
||
class SequencePacker:
|
||
"""序列打包（bin-packing）"""
|
||
|
||
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
|