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 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() 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