corentin/torch_utils
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add pipeline to SequenceBatchGenerator

Browse source 
* Remove flip_data parameter (can be done in the pipeline)
This commit is contained in:
Corentin 2021-03-23 02:13:40 +09:00
commit 73457750a8
2 changed files with 31 additions and 36 deletions
Download patch file Download diff file Expand all files Collapse all files

13
utils/batch_generator.py
View file

@ -68,10 +68,15 @@ class BatchGenerator:
self.global_step = 0
self.step = 0
first_data = np.array([data_processor(entry) if data_processor else entry
for entry in self.data[self.index_list[:batch_size]]])
first_label = np.array([label_processor(entry) if label_processor else entry
for entry in self.label[self.index_list[:batch_size]]])
first_data = [data_processor(entry) if data_processor else entry
for entry in self.data[self.index_list[:batch_size]]]
first_label = [label_processor(entry) if label_processor else entry
for entry in self.label[self.index_list[:batch_size]]]
if self.pipeline is not None:
for data_index, sample_data in enumerate(first_data):
first_data[data_index], first_label[data_index] = self.pipeline(sample_data, first_label[data_index])
first_data = np.asarray(first_data)
first_label = np.asarray(first_label)
self.batch_data = first_data
self.batch_label = first_label