def test_single_record(): channel, directory = write_to_channel("A", [b"bear"]) with tf.Session() as sess: dataset = PipeModeDataset(channel, pipe_dir=directory, state_dir=directory, config_dir=directory) it = dataset.make_one_shot_iterator() next = it.get_next() assert b"bear" == sess.run(next)
def test_multiple_iterators(): channel, directory = write_to_channel("A", [b"bear"]) dataset = PipeModeDataset(channel, pipe_dir=directory, state_dir=directory, config_dir=directory) with tf.Session() as sess: it = dataset.make_one_shot_iterator() next = it.get_next() assert sess.run(next) == b"bear" with pytest.raises(tf.errors.OutOfRangeError): sess.run(next) with open(os.path.join(directory, channel + "_1"), 'wb') as f: write_recordio(f, b"bunny") write_recordio(f, b"piano") write_recordio(f, b"caterpillar") with tf.Session() as sess: it = dataset.make_one_shot_iterator() next = it.get_next() assert b"bunny" == sess.run(next) assert b"piano" == sess.run(next) assert b"caterpillar" == sess.run(next) with pytest.raises(tf.errors.OutOfRangeError): sess.run(next)
def test_csv(): channel_dir = tempfile.mkdtemp() state_dir = tempfile.mkdtemp() epochs = 1 channel_name = 'testchannel' write_config(channel_dir, 'testchannel') create_fifos(epochs, channel_dir, channel_name, input_file='test.csv') def parse(line): fields = tf.decode_csv(line, FIELD_DEFAULTS) features = dict(zip(COLUMNS, fields)) return features with tf.Session() as sess: ds = PipeModeDataset(channel_name, pipe_dir=channel_dir, state_dir=state_dir, config_dir=channel_dir, record_format='TextLine') ds = ds.map(parse) it = ds.make_one_shot_iterator() next = it.get_next() for i in range(100): d = sess.run(next) sys.stdout.flush() assert d == {str(i): i for i in range(100)}
def input_fn(params): """The actual input function.""" batch_size = params['batch_size'] # For training, we want a lot of parallel reading and shuffling. # For eval, we want no shuffling and parallel reading doesn't matter. if is_training: channel = 'train' else: channel = 'validation' if pipe_mode: print('***** Using pipe_mode!!!!') ds = PipeModeDataset(channel=channel, record_format='TFRecord') else: ds = tf.data.TFRecordDataset(input_files) if is_training: ds = ds.repeat() ds = ds.shuffle(buffer_size=100) ds = ds.apply( tf.contrib.data.map_and_batch( lambda record: _decode_record(record, name_to_features), batch_size=batch_size, drop_remainder=drop_remainder)) return ds
def load_tfrecords_pipemode_dataset(self): """ STREAM dataset saved under the path "pipemode_SUBSET". Assumes a specific structure of tf records. Namely: "X" - flattened 1-dim length 60 "y" - 1-dim length 20 "seq_lens" - scalar Parameters ---------- example_proto Pointer to a TFRecord Returns ------- tf.data.Dataset-like SageMaker PipeModeDataset Pointer to k:v dictionary where value is the tensor. Keys: X, y, seq_lens """ if self.subset in ["train", "validation", "test"]: dataset = PipeModeDataset(channel="pipemode_" + self.subset, record_format="TFRecord") # Process individual samples from tfrecords to rank 1-3 tensors return dataset.map( lambda x: self.example_to_features(x, self.keys_to_features)) else: raise ValueError("Invalid data subset '%s'" % self.subset)
def test_multiple_records(): channel, directory = write_to_channel("B", [b"bunny", b"caterpillar"]) with tf.Session() as sess: dataset = PipeModeDataset(channel, pipe_dir=directory, state_dir=directory, config_dir=directory) it = dataset.make_one_shot_iterator() next = it.get_next() assert b"bunny" == sess.run(next) assert b"caterpillar" == sess.run(next)
def test_benchmark_records_interval_disabled(capfd): channel, directory = write_to_channel("A", [b"bear"]) dataset = PipeModeDataset(channel, pipe_dir=directory, state_dir=directory, config_dir=directory, benchmark_records_interval=0) with tf.Session() as sess: it = dataset.make_one_shot_iterator() next = it.get_next() assert sess.run(next) == b"bear" out, err = capfd.readouterr() assert 'Iterator records' not in out
def test_out_of_range(): channel, directory = write_to_channel("A", [b"bear", b"bunny", b"truck"]) with tf.Session() as sess: dataset = PipeModeDataset(channel, pipe_dir=directory, state_dir=directory, config_dir=directory) it = dataset.make_one_shot_iterator() next = it.get_next() for i in range(3): sess.run(next) with pytest.raises(tf.errors.OutOfRangeError): sess.run(next)
def test_large_record(): channel, directory = write_to_channel("C", [b"a" * 1000000]) with tf.Session() as sess: dataset = PipeModeDataset(channel, pipe_dir=directory, state_dir=directory) it = dataset.make_one_shot_iterator() next = it.get_next() assert b"a" * 1000000 == sess.run(next)
def get_dataset(batch_size, channel_name, dataset_bottleneck = False): from sagemaker_tensorflow import PipeModeDataset dataset = PipeModeDataset(channel_name, record_format='TFRecord').repeat() dataset = dataset.map(parse_image_function, num_parallel_calls=tf.data.experimental.AUTOTUNE) if dataset_bottleneck: dataset = dataset.map(data_augmentation, num_parallel_calls=tf.data.experimental.AUTOTUNE) dataset = dataset.batch(batch_size).prefetch(1) return dataset
def test_invalid_data(): directory = tempfile.mkdtemp() filename = "X_0" with open(directory + "/" + filename, 'wb') as f: f.write(b"adfsafasfd") write_config(directory, 'X') dataset = PipeModeDataset("X", pipe_dir=directory, state_dir=directory, config_dir=directory) with pytest.raises(tf.errors.InternalError): with tf.Session() as sess: it = dataset.make_one_shot_iterator() next = it.get_next() sess.run(next)
def process_input(epochs, batch_size, channel, channel_name, data_config): mode = data_config[channel_name]['TrainingInputMode'] filenames = _get_filenames(channel_name, channel) # Repeat infinitely. logging.info("Running {} in {} mode".format(channel_name, mode)) if mode == 'Pipe': from sagemaker_tensorflow import PipeModeDataset dataset = PipeModeDataset(channel=channel_name, record_format='TFRecord') else: dataset = tf.data.TFRecordDataset(filenames) dataset = dataset.repeat(epochs) dataset = dataset.prefetch(10) # Parse records. dataset = dataset.map(_dataset_parser, num_parallel_calls=10) # Potentially shuffle records. if channel_name == 'train': # Ensure that the capacity is sufficiently large to provide good random # shuffling. buffer_size = int( NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN * 0.4) + 3 * batch_size dataset = dataset.shuffle(buffer_size=buffer_size) # Batch it up. dataset = dataset.batch(batch_size, drop_remainder=True) iterator = dataset.make_one_shot_iterator() image_batch, label_batch = iterator.get_next() return image_batch, label_batch
def _input(epochs, batch_size, channel, channel_name, hvd=None): if hvd != None: channel_name = '{}_{}'.format(channel_name, hvd.rank() % 4) print("The channel name is ", channel_name) channel_input_dir = args.training_env['channel_input_dirs'][channel_name] print("The corresponding input directory is ", channel_input_dir) mode = args.data_config[channel_name]['TrainingInputMode'] if mode == 'Pipe': from sagemaker_tensorflow import PipeModeDataset dataset = PipeModeDataset(channel=channel_name, record_format='TFRecord') else: filenames = get_filenames(channel_input_dir, hvd) print("The correpsonding filenames are", filenames) dataset = tf.data.TFRecordDataset(filenames) if 'test' in channel_name: dataset = dataset.map(_dataset_parser_with_slide) else: dataset = dataset.repeat(epochs) dataset = dataset.map(_dataset_parser) if 'train' in channel_name: # Ensure that the capacity is sufficiently large to provide good random shuffling. buffer_size = int( NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN * 0.4) + 3 * batch_size dataset = dataset.shuffle(buffer_size=buffer_size) # Batch it up (only for train and valid) if 'test' not in channel_name: dataset = dataset.batch(batch_size, drop_remainder=True) dataset = dataset.prefetch(10) return dataset
def _input(epochs, batch_size, channel, channel_name): filenames = get_filenames(channel_name, channel) # ----- 추가 부분 (PipeModeDataSet) ----- #dataset = tf.data.TFRecordDataset(filenames) dataset = PipeModeDataset(channel=channel_name, record_format='TFRecord') dataset = dataset.repeat(epochs) dataset = dataset.prefetch(10) # Parse records. dataset = dataset.map( _dataset_parser, num_parallel_calls=10) # Potentially shuffle records. if channel_name == 'train': # Ensure that the capacity is sufficiently large to provide good random # shuffling. buffer_size = int(NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN * 0.4) + 3 * batch_size dataset = dataset.shuffle(buffer_size=buffer_size) # Batch it up. dataset = dataset.batch(batch_size, drop_remainder=True) iterator = dataset.make_one_shot_iterator() image_batch, label_batch = iterator.get_next() return {INPUT_TENSOR_NAME: image_batch}, label_batch
def _input_fn(): def _read_and_decode(record): features = tf.parse_single_example( record, features={ 'image_raw': tf.FixedLenFeature([], tf.string), 'label': tf.FixedLenFeature([], tf.int64), }) image = tf.decode_raw(features['image_raw'], tf.uint8) image.set_shape([HEIGHT * WIDTH * DEPTH]) image = tf.cast(image, tf.float32) * (1. / 255) label = tf.cast(features['label'], tf.int32) return {INPUT_TENSOR_NAME: image}, label ds = PipeModeDataset(channel, record_format='TFRecord') ds = ds.repeat() ds = ds.prefetch(batch_size) ds = ds.map(_read_and_decode, num_parallel_calls=NUM_PARALLEL_BATCHES) if channel == 'train': ds = ds.shuffle(buffer_size=batch_size) ds = ds.batch(batch_size, drop_remainder=True) ds = ds.make_one_shot_iterator().get_next() return ds
def test_input_config_validation_failure(): channel_dir = tempfile.mkdtemp() state_dir = tempfile.mkdtemp() write_config(channel_dir, 'testchannel') with pytest.raises(PipeModeDatasetException): with tf.Session(): PipeModeDataset("Not a Channel", pipe_dir=channel_dir, state_dir=state_dir, config_dir=channel_dir)
def test_tf_record(): channel_dir = tempfile.mkdtemp() state_dir = tempfile.mkdtemp() epochs = 1 channel_name = 'testchannel' create_fifos(epochs, channel_dir, channel_name, input_file='test.tfrecords') write_config(channel_dir, 'testchannel') ds = PipeModeDataset(channel_name, pipe_dir=channel_dir, state_dir=state_dir, config_dir=channel_dir, record_format='TFRecord') with tf.Session() as sess: it = ds.make_one_shot_iterator() next = it.get_next() for i in range(100): assert sess.run(next) == b'hello world'
def file_based_input_dataset_builder(channel, input_filenames, pipe_mode, is_training, drop_remainder, batch_size, epochs, steps_per_epoch, max_seq_length): # For training, we want a lot of parallel reading and shuffling. # For eval, we want no shuffling and parallel reading doesn't matter. if pipe_mode: print('***** Using pipe_mode with channel {}'.format(channel)) from sagemaker_tensorflow import PipeModeDataset dataset = PipeModeDataset(channel=channel, record_format='TFRecord') else: print('***** Using input_filenames {}'.format(input_filenames)) dataset = tf.data.TFRecordDataset(input_filenames) dataset = dataset.repeat(epochs * steps_per_epoch * 100) dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) name_to_features = { "input_ids": tf.io.FixedLenFeature([max_seq_length], tf.int64), "input_mask": tf.io.FixedLenFeature([max_seq_length], tf.int64), "segment_ids": tf.io.FixedLenFeature([max_seq_length], tf.int64), "label_ids": tf.io.FixedLenFeature([], tf.int64), } def _decode_record(record, name_to_features): """Decodes a record to a TensorFlow example.""" record = tf.io.parse_single_example(record, name_to_features) # TODO: wip/bert/bert_attention_head_view/train.py # Convert input_ids into input_tokens with DistilBert vocabulary # if hook.get_collections()['all'].save_config.should_save_step(modes.EVAL, hook.mode_steps[modes.EVAL]): # hook._write_raw_tensor_simple("input_tokens", input_tokens) return record dataset = dataset.apply( tf.data.experimental.map_and_batch( lambda record: _decode_record(record, name_to_features), batch_size=batch_size, drop_remainder=drop_remainder, num_parallel_calls=tf.data.experimental.AUTOTUNE)) dataset.cache() if is_training: dataset = dataset.shuffle(seed=42, buffer_size=100, reshuffle_each_iteration=True) return dataset
def test_missing_channel(): channel, directory = write_to_channel("A", [b"bear", b"bunny", b"truck"]) with tf.Session() as sess: with pytest.raises(PipeModeDatasetException): PipeModeDataset("Not A Channel", pipe_dir=directory, state_dir=directory, config_dir=directory)
def _input(epochs, batch_size, channel, channel_name): """Uses the tf.data input pipeline for CIFAR-10 dataset.""" mode = args.data_config[channel_name]["TrainingInputMode"] logging.info("Running {} in {} mode".format(channel_name, mode)) if mode == "Pipe": from sagemaker_tensorflow import PipeModeDataset dataset = PipeModeDataset(channel=channel_name, record_format="TFRecord") else: filenames = _get_filenames(channel_name, channel) dataset = tf.data.TFRecordDataset(filenames) # Repeat infinitely. dataset = dataset.repeat() dataset = dataset.prefetch(10) # Parse records. dataset = dataset.map(_dataset_parser, num_parallel_calls=10) # Potentially shuffle records. if channel_name == "train": # Ensure that the capacity is sufficiently large to provide good random shuffling. buffer_size = int( NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN * 0.4) + 3 * batch_size dataset = dataset.shuffle(buffer_size=buffer_size) # Batch it up. dataset = dataset.batch(batch_size, drop_remainder=True) iterator = tf.compat.v1.data.make_one_shot_iterator(dataset) image_batch, label_batch = iterator.get_next() return {INPUT_TENSOR_NAME: image_batch}, label_batch
def _input(epochs, batch_size, channel, channel_name): mode = args.data_config[channel_name]['TrainingInputMode'] """Uses the tf.data input pipeline for our dataset. Args: mode: Standard names for model modes (tf.estimators.ModeKeys). batch_size: The number of samples per batch of input requested. """ logging.info("Running {} in {} mode for {} epochs".format(channel_name, mode, epochs)) filenames = get_filenames(channel_name, channel) # Repeat infinitely. if mode == 'Pipe': from sagemaker_tensorflow import PipeModeDataset dataset = PipeModeDataset(channel=channel_name, record_format='TFRecord') else: dataset = tf.data.TFRecordDataset(filenames) dataset = dataset.repeat(epochs) dataset = dataset.prefetch(batch_size) # Parse records. dataset = dataset.map(_dataset_parser, num_parallel_calls=10) # Shuffle training records. if channel_name == 'train': # Ensure that the capacity is sufficiently large to provide good random # shuffling. buffer_size = args.num_train_samples // args.batch_size dataset = dataset.shuffle(buffer_size=buffer_size) # Batch it up. dataset = dataset.batch(batch_size, drop_remainder=True) if tf.version.VERSION[0] == '2': iterator = tf.compat.v1.data.make_one_shot_iterator(dataset) else: iterator = dataset.make_one_shot_iterator() features_batch, label_batch = iterator.get_next() if tf.version.VERSION[0] == '2': with tf.compat.v1.Session() as sess: logging.info('type of features_batch: {}, type of values: {}'.format(type(features_batch), type(features_batch))) logging.info('label_batch: {}'.format(label_batch)) logging.info('type of label_batch: {}'.format(type(label_batch))) else: with tf.Session() as sess: logging.info('type of features_batch: {}, type of values: {}'.format(type(features_batch), type(features_batch))) logging.info('label_batch: {}'.format(label_batch)) logging.info('type of label_batch: {}'.format(type(label_batch))) return {INPUT_TENSOR_NAME: features_batch}, label_batch
def _input_fn(channel): """Returns a Dataset for reading from a SageMaker PipeMode channel.""" features = { 'label': tf.FixedLenFeature([], tf.int64), 'feature': tf.FixedLenFeature([480], tf.int64) } def parse(record): parsed = tf.parse_single_example(record, features) data = tf.reshape(parsed['feature'], [160, 3]) data = tf.cast(data, tf.float32) label = tf.cast(parsed['label'], tf.int32) return ({INPUT_TENSOR_NAME: data}, label) ds = PipeModeDataset(channel=channel, record_format='TFRecord') ds = ds.shuffle(SHUFFLE_BUFFER_SIZE) ds = ds.repeat(MAX_EPOCHS) ds = ds.prefetch(PREFETCH_SIZE) ds = ds.map(parse, num_parallel_calls=NUM_PARALLEL_BATCHES) ds = ds.batch(BATCH_SIZE) return ds
def input_fn(): ds = PipeModeDataset(channel_name, pipe_dir=channel_dir, state_dir=state_dir, config_dir=channel_dir) ds = ds.map(parse, num_parallel_calls=12) ds = ds.repeat(count=2) ds = ds.prefetch(3) ds = ds.batch(3) it = ds.make_one_shot_iterator() return it.get_next()
def _input_fn(): features = { 'data': tf.FixedLenFeature([], tf.string), 'labels': tf.FixedLenFeature([], tf.int64), } def parse(record): return tf.parse_single_example(record, features) ds = PipeModeDataset(config.channel, benchmark=True) if config.epochs > 1: ds = ds.repeat(config.epochs) if config.prefetch_size > 0: ds = ds.prefetch(config.prefetch_size) ds = ds.apply( map_and_batch(parse, batch_size=config.batch_size, num_parallel_batches=config.parallel_transform_calls)) return ds
def _input_fn(channel): """Returns a Dataset for reading from a SageMaker PipeMode channel.""" features = { 'data': tf.FixedLenFeature([], tf.string), 'labels': tf.FixedLenFeature([], tf.int64), } def parse(record): parsed = tf.parse_single_example(record, features) return ({ 'data': tf.decode_raw(parsed['data'], tf.float64) }, parsed['labels']) ds = PipeModeDataset(channel) if EPOCHS > 1: ds = ds.repeat(EPOCHS) ds = ds.prefetch(PREFETCH_SIZE) ds = ds.apply(map_and_batch(parse, batch_size=BATCH_SIZE, num_parallel_batches=NUM_PARALLEL_BATCHES)) return ds
def _input_fn(channel): """Returns a Dataset which reads from a SageMaker PipeMode channel.""" features = { 'image_raw': tf.FixedLenFeature([], tf.string), 'label': tf.FixedLenFeature([], tf.int64), 'height': tf.FixedLenFeature([], tf.int64), 'width': tf.FixedLenFeature([], tf.int64), 'channels': tf.FixedLenFeature([], tf.int64) } def parse(record): parsed = tf.parse_single_example(record, features) image = tf.decode_raw(parsed['image_raw'], tf.uint8) image.set_shape([784]) image = tf.cast(image, tf.float32) * (1. / 255) label = tf.cast(parsed['label'], tf.int32) return ({INPUT_TENSOR_NAME: image}, label) ds = PipeModeDataset(channel=channel, record_format='TFRecord') ds = ds.repeat(MAX_EPOCHS) ds = ds.prefetch(PREFETCH_SIZE) ds = ds.map(parse, num_parallel_calls=NUM_PARALLEL_BATCHES) ds = ds.batch(BATCH_SIZE) return ds
def make_dataset(channel_name): ds = PipeModeDataset(channel_name, pipe_dir=channel_dir, state_dir=state_dir, config_dir=channel_dir) ds = ds.map(parse, num_parallel_calls=12) ds = ds.repeat(count=2) ds = ds.prefetch(3) ds = ds.batch(10) return ds
def file_based_input_dataset_builder( channel, input_filenames, pipe_mode, is_training, drop_remainder, batch_size, epochs, steps_per_epoch, max_seq_length, ): # For training, we want a lot of parallel reading and shuffling. # For eval, we want no shuffling and parallel reading doesn't matter. if pipe_mode: print("***** Using pipe_mode with channel {}".format(channel)) from sagemaker_tensorflow import PipeModeDataset dataset = PipeModeDataset(channel=channel, record_format="TFRecord") else: print("***** Using input_filenames {}".format(input_filenames)) dataset = tf.data.TFRecordDataset(input_filenames) dataset = dataset.repeat(epochs * steps_per_epoch * 100) name_to_features = { "input_ids": tf.io.FixedLenFeature([max_seq_length], tf.int64), "input_mask": tf.io.FixedLenFeature([max_seq_length], tf.int64), "segment_ids": tf.io.FixedLenFeature([max_seq_length], tf.int64), "label_ids": tf.io.FixedLenFeature([], tf.int64), } def _decode_record(record, name_to_features): """Decodes a record to a TensorFlow example.""" record = tf.io.parse_single_example(record, name_to_features) return record dataset = dataset.apply( tf.data.experimental.map_and_batch( lambda record: _decode_record(record, name_to_features), batch_size=batch_size, drop_remainder=drop_remainder, num_parallel_calls=tf.data.experimental.AUTOTUNE, ) ) dataset = dataset.shuffle(buffer_size=1000, reshuffle_each_iteration=True) row_count = 0 print("**************** {} *****************".format(channel)) for row in dataset.as_numpy_iterator(): print(row) if row_count == 5: break row_count = row_count + 1 return dataset
def _input_fn(channel): """Returns a Dataset for reading from a SageMaker PipeMode channel.""" features = { "data": tf.FixedLenFeature([], tf.string), "labels": tf.FixedLenFeature([], tf.int64), } def parse(record): parsed = tf.parse_single_example(record, features) return ({ "data": tf.decode_raw(parsed["data"], tf.float64) }, parsed["labels"]) ds = PipeModeDataset(channel) if EPOCHS > 1: ds = ds.repeat(EPOCHS) ds = ds.prefetch(PREFETCH_SIZE) ds = ds.apply( map_and_batch(parse, batch_size=BATCH_SIZE, num_parallel_batches=NUM_PARALLEL_BATCHES)) return ds
def _input(epochs, batch_size, channel, channel_name, hvd=None): # If Horovod, assign channel name using the horovod rank if hvd != None: channel_name = '{}_{}'.format(channel_name, hvd.local_rank()) channel_input_dir = args.training_env['channel_input_dirs'][channel_name] mode = args.data_config[channel_name]['TrainingInputMode'] """Uses the tf.data input pipeline for CIFAR-10 dataset. Args: mode: Standard names for model modes (tf.estimators.ModeKeys). batch_size: The number of samples per batch of input requested. """ if mode == 'Pipe': from sagemaker_tensorflow import PipeModeDataset dataset = PipeModeDataset(channel=channel_name, record_format='TFRecord') #, benchmark=True) else: filenames = get_filenames(channel_input_dir) print(f'DEBUG tfrecords : {filenames}') dataset = tf.data.TFRecordDataset(filenames) if 'train' in channel_name: dataset = dataset.repeat(epochs) else: dataset = dataset.repeat(20) # Parse records. dataset = dataset.map(_dataset_parser, num_parallel_calls=10) # Potentially shuffle records. # if hvd == None and 'train' in channel_name: if 'train' in channel_name: # Ensure that the capacity is sufficiently large to provide good random # shuffling. buffer_size = int( NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN * 0.4) + 3 * batch_size dataset = dataset.shuffle(buffer_size=buffer_size) # Batch it up. dataset = dataset.batch(batch_size, drop_remainder=True) dataset = dataset.prefetch(10) return dataset
def _input(args, channel_name): try: mode_channel_name = channel_name + 'ing' if channel_name == 'train' else channel_name mode = args.data_config[mode_channel_name]['TrainingInputMode'] except: mode = 'File' """Uses the tf.data input pipeline for dataset. Args: mode: Standard names for model modes (tf.estimators.ModeKeys). batch_size: The number of samples per batch of input requested. """ filenames = get_filenames(args, channel_name) # Repeat infinitely. logging.info("Running {} in {} mode".format(channel_name, mode)) if mode == 'Pipe': from sagemaker_tensorflow import PipeModeDataset dataset = PipeModeDataset(channel=channel_name, record_format='TFRecord') else: dataset = tf.data.TFRecordDataset(filenames) # Potentially shuffle records. if channel_name == 'train': # Ensure that the capacity is sufficiently large to provide good random # shuffling. dataset = dataset.map(_load_image_train, num_parallel_calls=tf.data.experimental.AUTOTUNE) buffer_size = int(args.train_num_examples * 0.4) + 3 * args.BATCH_SIZE dataset = dataset.cache().shuffle(buffer_size=buffer_size).batch( args.BATCH_SIZE).repeat() elif channel_name == 'test': dataset = dataset.map(_load_image_test) for image, mask in dataset.take(1): sample_image, sample_mask = image, mask _img_save('sample_image.jpg', sample_image) _img_save('sample_mask.png', sample_mask) dataset = dataset.batch(args.BATCH_SIZE) dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE) return dataset