def input_fn(data_dir,
subset,
batch_size,
use_distortion_for_training=True,
shuffle=False):
use_distortion = subset == 'train' and use_distortion_for_training
dataset = dataset_utils.ConvLSTMDataSet(data_dir, subset, use_distortion)
return dataset.make_batch(batch_size, shuffle)
def input_fn_multigpus(data_dir, subset, batch_size,
use_distortion_for_training=True):
use_distortion = subset == 'train' and use_distortion_for_training
dataset = dataset_utils.ConvLSTMDataSet(data_dir, subset, use_distortion)
dataset = dataset.make_batch(batch_size, True)
return dataset.prefetch(NUM_GPUS)