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 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()
if is_training:
dataset = dataset.shuffle(seed=42,
buffer_size=100,
reshuffle_each_iteration=True)
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)
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),
# "is_real_example": tf.io.FixedLenFeature([], tf.int64),
}
def _decode_record(record, name_to_features):
"""Decodes a record to a TensorFlow example."""
return tf.io.parse_single_example(record, name_to_features)
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=1000,
reshuffle_each_iteration=True)
return dataset
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 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(filenames,
channel='training',
batch_size=32,
num_epochs=1,
perform_shuffle=False):
print('Parsing', filenames)
def decode_libsvm(line):
#columns = tf.decode_csv(value, record_defaults=CSV_COLUMN_DEFAULTS)
#features = dict(zip(CSV_COLUMNS, columns))
#labels = features.pop(LABEL_COLUMN)
columns = tf.string_split([line], ' ')
labels = tf.string_to_number(columns.values[0], out_type=tf.float32)
splits = tf.string_split(columns.values[1:], ':')
id_vals = tf.reshape(splits.values, splits.dense_shape)
feat_ids, feat_vals = tf.split(id_vals, num_or_size_splits=2, axis=1)
feat_ids = tf.string_to_number(feat_ids, out_type=tf.int32)
feat_vals = tf.string_to_number(feat_vals, out_type=tf.float32)
#feat_ids = tf.reshape(feat_ids,shape=[-1,FLAGS.field_size])
#for i in range(splits.dense_shape.eval()[0]):
# feat_ids.append(tf.string_to_number(splits.values[2*i], out_type=tf.int32))
# feat_vals.append(tf.string_to_number(splits.values[2*i+1]))
#return tf.reshape(feat_ids,shape=[-1,field_size]), tf.reshape(feat_vals,shape=[-1,field_size]), labels
return {"feat_ids": feat_ids, "feat_vals": feat_vals}, labels
# Extract lines from input files using the Dataset API, can pass one filename or filename list
if FLAGS.pipe_mode == 0:
dataset = tf.data.TextLineDataset(filenames).map(
decode_libsvm, num_parallel_calls=10).prefetch(
500000) # multi-thread pre-process then prefetch
# Randomizes input using a window of 256 elements (read into memory)
if perform_shuffle:
dataset = dataset.shuffle(buffer_size=256)
# epochs from blending together.
dataset = dataset.repeat(num_epochs)
#liangaws:注意如果是单机多GPU或者多CPU，这里的batch_size应设置为CPU或者GPU数量的倍数来充分利用算力。
#dataset = dataset.batch(batch_size) # Batch size to use
#liangaws:这里使用drop_remainder=True来把不够一个batch size的数据忽略
dataset = dataset.batch(batch_size,
drop_remainder=True) # Batch size to use
"""
#return dataset.make_one_shot_iterator()
iterator = dataset.make_one_shot_iterator()
batch_features, batch_labels = iterator.get_next()
#return tf.reshape(batch_ids,shape=[-1,field_size]), tf.reshape(batch_vals,shape=[-1,field_size]), batch_labels
return batch_features, batch_labels
"""
#liangaws:当使用tensorflow的dataset API与distribute strategy联合使用的时候，input_fn需要返回dataset而不是返回特征和label
return dataset
else:
print("-------enter into pipe mode branch!------------")
dataset = PipeModeDataset(channel, record_format='TextLine')
if num_epochs > 1:
dataset = dataset.repeat(num_epochs)
dataset = dataset.prefetch(batch_size * 100)
dataset = dataset.apply(
map_and_batch(decode_libsvm,
batch_size=batch_size,
num_parallel_batches=10))
return dataset