From 51776d699941fb2d1811773a36ce35a9e2f6c30e Mon Sep 17 00:00:00 2001 From: Corentin Date: Thu, 18 Mar 2021 10:36:17 +0900 Subject: [PATCH] Sequence stride option for SequenceBatchGenerator --- utils/sequence_batch_generator.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/utils/sequence_batch_generator.py b/utils/sequence_batch_generator.py index 04cd55e..041bd60 100644 --- a/utils/sequence_batch_generator.py +++ b/utils/sequence_batch_generator.py @@ -17,7 +17,7 @@ class SequenceGenerator(BatchGenerator): def __init__(self, data: Iterable, label: Iterable, sequence_size: int, batch_size: int, data_processor: Optional[Callable] = None, label_processor: Optional[Callable] = None, prefetch=True, preload=False, num_workers=1, shuffle=True, initial_shuffle=False, - flip_data=False, save: Optional[str] = None): + sequence_stride=1, flip_data=False, save: Optional[str] = None): self.batch_size = batch_size self.sequence_size = sequence_size self.shuffle = shuffle @@ -72,9 +72,14 @@ class SequenceGenerator(BatchGenerator): self.index_list = [] for sequence_index in range(len(self.data)): - start_indices = np.expand_dims( - np.arange(len(self.data[sequence_index]) - sequence_size + 1, dtype=np.uint32), - axis=-1) + if sequence_stride > 1: + start_indices = np.expand_dims( + np.arange(0, len(self.data[sequence_index]) - sequence_size + 1, sequence_stride, dtype=np.uint32), + axis=-1) + else: + start_indices = np.expand_dims( + np.arange(len(self.data[sequence_index]) - sequence_size + 1, dtype=np.uint32), + axis=-1) start_indices = np.insert(start_indices, 0, sequence_index, axis=1) self.index_list.append(start_indices) self.index_list = np.concatenate(self.index_list, axis=0)