File size: 6,206 Bytes
ad16788 |
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 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
from typing import List
from typing import Sequence
from typing import Tuple
from typing import Union
from typeguard import check_argument_types
from typeguard import check_return_type
from espnet2.samplers.abs_sampler import AbsSampler
from espnet2.samplers.folded_batch_sampler import FoldedBatchSampler
from espnet2.samplers.length_batch_sampler import LengthBatchSampler
from espnet2.samplers.num_elements_batch_sampler import NumElementsBatchSampler
from espnet2.samplers.sorted_batch_sampler import SortedBatchSampler
from espnet2.samplers.unsorted_batch_sampler import UnsortedBatchSampler
BATCH_TYPES = dict(
unsorted="UnsortedBatchSampler has nothing in paticular feature and "
"just creates mini-batches which has constant batch_size. "
"This sampler doesn't require any length "
"information for each feature. "
"'key_file' is just a text file which describes each sample name."
"\n\n"
" utterance_id_a\n"
" utterance_id_b\n"
" utterance_id_c\n"
"\n"
"The fist column is referred, so 'shape file' can be used, too.\n\n"
" utterance_id_a 100,80\n"
" utterance_id_b 400,80\n"
" utterance_id_c 512,80\n",
sorted="SortedBatchSampler sorts samples by the length of the first input "
" in order to make each sample in a mini-batch has close length. "
"This sampler requires a text file which describes the length for each sample "
"\n\n"
" utterance_id_a 1000\n"
" utterance_id_b 1453\n"
" utterance_id_c 1241\n"
"\n"
"The first element of feature dimensions is referred, "
"so 'shape_file' can be also used.\n\n"
" utterance_id_a 1000,80\n"
" utterance_id_b 1453,80\n"
" utterance_id_c 1241,80\n",
folded="FoldedBatchSampler supports variable batch_size. "
"The batch_size is decided by\n"
" batch_size = base_batch_size // (L // fold_length)\n"
"L is referred to the largest length of samples in the mini-batch. "
"This samples requires length information as same as SortedBatchSampler\n",
length="LengthBatchSampler supports variable batch_size. "
"This sampler makes mini-batches which have same number of 'bins' as possible "
"counting by the total lengths of each feature in the mini-batch. "
"This sampler requires a text file which describes the length for each sample. "
"\n\n"
" utterance_id_a 1000\n"
" utterance_id_b 1453\n"
" utterance_id_c 1241\n"
"\n"
"The first element of feature dimensions is referred, "
"so 'shape_file' can be also used.\n\n"
" utterance_id_a 1000,80\n"
" utterance_id_b 1453,80\n"
" utterance_id_c 1241,80\n",
numel="NumElementsBatchSampler supports variable batch_size. "
"Just like LengthBatchSampler, this sampler makes mini-batches"
" which have same number of 'bins' as possible "
"counting by the total number of elements of each feature "
"instead of the length. "
"Thus this sampler requires the full information of the dimension of the features. "
"\n\n"
" utterance_id_a 1000,80\n"
" utterance_id_b 1453,80\n"
" utterance_id_c 1241,80\n",
)
def build_batch_sampler(
type: str,
batch_size: int,
batch_bins: int,
shape_files: Union[Tuple[str, ...], List[str]],
sort_in_batch: str = "descending",
sort_batch: str = "ascending",
drop_last: bool = False,
min_batch_size: int = 1,
fold_lengths: Sequence[int] = (),
padding: bool = True,
utt2category_file: str = None,
) -> AbsSampler:
"""Helper function to instantiate BatchSampler.
Args:
type: mini-batch type. "unsorted", "sorted", "folded", "numel", or, "length"
batch_size: The mini-batch size. Used for "unsorted", "sorted", "folded" mode
batch_bins: Used for "numel" model
shape_files: Text files describing the length and dimension
of each features. e.g. uttA 1330,80
sort_in_batch:
sort_batch:
drop_last:
min_batch_size: Used for "numel" or "folded" mode
fold_lengths: Used for "folded" mode
padding: Whether sequences are input as a padded tensor or not.
used for "numel" mode
"""
assert check_argument_types()
if len(shape_files) == 0:
raise ValueError("No shape file are given")
if type == "unsorted":
retval = UnsortedBatchSampler(
batch_size=batch_size, key_file=shape_files[0], drop_last=drop_last
)
elif type == "sorted":
retval = SortedBatchSampler(
batch_size=batch_size,
shape_file=shape_files[0],
sort_in_batch=sort_in_batch,
sort_batch=sort_batch,
drop_last=drop_last,
)
elif type == "folded":
if len(fold_lengths) != len(shape_files):
raise ValueError(
f"The number of fold_lengths must be equal to "
f"the number of shape_files: "
f"{len(fold_lengths)} != {len(shape_files)}"
)
retval = FoldedBatchSampler(
batch_size=batch_size,
shape_files=shape_files,
fold_lengths=fold_lengths,
sort_in_batch=sort_in_batch,
sort_batch=sort_batch,
drop_last=drop_last,
min_batch_size=min_batch_size,
utt2category_file=utt2category_file,
)
elif type == "numel":
retval = NumElementsBatchSampler(
batch_bins=batch_bins,
shape_files=shape_files,
sort_in_batch=sort_in_batch,
sort_batch=sort_batch,
drop_last=drop_last,
padding=padding,
min_batch_size=min_batch_size,
)
elif type == "length":
retval = LengthBatchSampler(
batch_bins=batch_bins,
shape_files=shape_files,
sort_in_batch=sort_in_batch,
sort_batch=sort_batch,
drop_last=drop_last,
padding=padding,
min_batch_size=min_batch_size,
)
else:
raise ValueError(f"Not supported: {type}")
assert check_return_type(retval)
return retval
|