File size: 4,797 Bytes
6b73a07 d572bbd 6b73a07 d572bbd 6b73a07 d572bbd 6b73a07 d572bbd 6b73a07 d572bbd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """Batch collation and dynamic batching for translation training."""
from __future__ import annotations
import random
from typing import Iterator, Sequence
import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Sampler
class TranslationCollator:
"""Pad variable-length translation examples into one batch."""
def __init__(self, pad_token_id: int = 0, label_pad_token_id: int = -100):
self.pad_token_id = pad_token_id
self.label_pad_token_id = label_pad_token_id
def __call__(self, batch: Sequence[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
src_ids = pad_sequence(
[item["src_ids"] for item in batch],
batch_first=True,
padding_value=self.pad_token_id,
)
tgt_input_ids = pad_sequence(
[item["tgt_input_ids"] for item in batch],
batch_first=True,
padding_value=self.pad_token_id,
)
labels = pad_sequence(
[item["labels"] for item in batch],
batch_first=True,
padding_value=self.label_pad_token_id,
)
src_padding_mask = src_ids.eq(self.pad_token_id)
tgt_padding_mask = tgt_input_ids.eq(self.pad_token_id)
return {
"src_ids": src_ids,
"tgt_input_ids": tgt_input_ids,
"labels": labels,
"src_padding_mask": src_padding_mask,
"tgt_padding_mask": tgt_padding_mask,
"src_attention_mask": (~src_padding_mask).long(),
"tgt_attention_mask": (~tgt_padding_mask).long(),
"src_lens": torch.stack([item["src_len"] for item in batch]),
"tgt_lens": torch.stack([item["tgt_len"] for item in batch]),
}
class DynamicBatchSampler(Sampler[list[int]]):
"""Create batches constrained by an approximate max token budget."""
def __init__(
self,
lengths: Sequence[int | tuple[int, int]],
max_tokens_per_batch: int = 8192,
shuffle: bool = True,
drop_last: bool = False,
):
self.lengths = [max(length) if isinstance(length, tuple) else int(length) for length in lengths]
self.max_tokens_per_batch = max_tokens_per_batch
self.shuffle = shuffle
self.drop_last = drop_last
# 检查是否有超过最大token限制的序列
self._check_long_sequences()
def _check_long_sequences(self):
"""检查并警告过长的序列"""
long_seq_count = sum(1 for l in self.lengths if l > self.max_tokens_per_batch)
if long_seq_count > 0:
import warnings
warnings.warn(
f"Found {long_seq_count} sequences longer than max_tokens_per_batch "
f"({self.max_tokens_per_batch}). These will be placed in their own batches."
)
def __iter__(self) -> Iterator[list[int]]:
indices = list(range(len(self.lengths)))
if self.shuffle:
random.shuffle(indices)
indices.sort(key=lambda idx: self.lengths[idx])
batches: list[list[int]] = []
batch: list[int] = []
max_len = 0
for idx in indices:
seq_len = self.lengths[idx]
# 如果单个序列长度超过限制，单独放入一个批次
if seq_len > self.max_tokens_per_batch:
if batch:
batches.append(batch)
batch = []
max_len = 0
batches.append([idx])
continue
candidate_max_len = max(max_len, seq_len)
candidate_tokens = candidate_max_len * (len(batch) + 1)
if batch and candidate_tokens > self.max_tokens_per_batch:
batches.append(batch)
batch = []
max_len = 0
batch.append(idx)
max_len = max(max_len, seq_len)
if batch and not self.drop_last:
batches.append(batch)
if self.shuffle:
random.shuffle(batches)
yield from batches
def __len__(self) -> int:
count = 0
batch_size = 0
max_len = 0
for length in sorted(self.lengths):
# 过长的序列单独计数
if length > self.max_tokens_per_batch:
count += 1
continue
candidate_max_len = max(max_len, length)
if batch_size and candidate_max_len * (batch_size + 1) > self.max_tokens_per_batch:
count += 1
batch_size = 0
max_len = 0
batch_size += 1
max_len = max(max_len, length)
if batch_size and not self.drop_last:
count += 1
return count |