【云+社区年度征文】tensorflow2 tfrecorddataset+estimator 训练预测加载全流程概述

MULTI_COLUMNS = []
SINGLE_COLUMNS = []
NUMBER_COLUMNS = []

writer = tf.io.TFRecordWriter(write_url)

for row in dataframe.itertuples():
    feature_map = {}
    feature_map['label'] = tf.train.Feature(int64_list=tf.train.Int64List(value=[getattr(row,'label')]))
    
    for i in MULTI_COLUMNS:
        feature_map[i] = tf.train.Feature(bytes_list=tf.train.BytesList(value=split(getattr(row,i))))
    for j in SINGLE_COLUMNS:
        feature_map[j] = tf.train.Feature(bytes_list=tf.train.BytesList(value=[bytes(getattr(row,j) , encoding = 'utf-8')]))
    for l in NUMBER_COLUMNS:
        feature_map[l] = tf.train.Feature(float_list=tf.train.FloatList(value=[getattr(row,l)]))

    example = tf.train.Example(
        features=tf.train.Features(
            feature = feature_map
        )
    )
    writer.write(record=example.SerializeToString())
writer.close()

def train_input_fn():
    train_ds = tf.data.TFRecordDataset(filenames=train_data_dir,num_parallel_reads = 104).shuffle(SUFFLE_SIZE, reshuffle_each_iteration=True).batch(65536).prefetch(4).repeat()
    return train_ds

def eval_input_fn():
    test_ds = tf.data.TFRecordDataset(filenames=eval_data_dir,num_parallel_reads = 104).shuffle(SUFFLE_SIZE, reshuffle_each_iteration=True).batch(65536).prefetch(4).repeat()
    return test_ds

all = tf.data.Dataset.from_tensor_slices(list(range(1, 21))) \\
        .shuffle(10, reshuffle_each_iteration=False)

test_dataset = all.enumerate() \\
                    .filter(lambda x,y: x % 4 == 0) \\
                    .map(lambda x,y: y)

train_dataset = all.enumerate() \\
                    .filter(lambda x,y: x % 4 != 0) \\
                    .map(lambda x,y: y)

tf.feature_column.numeric_column(
    key,
    shape=(1,),
    default_value=None,
    dtype=tf.float32,
    normalizer_fn=None
)

tf.feature_column.bucketized_column(
    source_column,
    boundaries
)

# First, convert the raw input to a numeric column.
numeric_feature_column = tf.feature_column.numeric_column("Year")

# Then, bucketize the numeric column on the years 1960, 1980, and 2000.
bucketized_feature_column = tf.feature_column.bucketized_column(
    source_column = numeric_feature_column,
    boundaries = [1960, 1980, 2000])

identity_feature_column = tf.feature_column.categorical_column_with_identity(
    key='my_feature_b',
    num_buckets=4) # Values [0, 4)

tf.feature_column.crossed_column(
    keys,
    hash_bucket_size,
    hash_key=None
)

tf.feature_column.embedding_column(
    categorical_column,
    dimension,
    combiner='mean',
    initializer=None,
    ckpt_to_load_from=None,
    tensor_name_in_ckpt=None,
    max_norm=None,
    trainable=True
)
categorical_column: 使用categoryical_column产生的sparsor column
dimension: 定义embedding的维数
combiner: 对于多个entries进行的推导。默认是meam, 但是 sqrtn在词袋模型中，有更好的准确度。
initializer: 初始化方法，默认使用高斯分布来初始化。
tensor_name_in_ckpt: 可以从check point中恢复
ckpt_to_load_from: check point file，这是在 tensor_name_in_ckpt 不为空的情况下设置的.
max_norm: 默认是l2

tf.feature_column.weighted_categorical_column(
    categorical_column,
    weight_feature_key,
    dtype=tf.float32
)

hidden_units = [256, 64, 32]
dropout = 0.3
my_dnn_optimizer = tf.keras.optimizers.Nadam(learning_rate=0.001)
config = tf.estimator.RunConfig(log_step_count_steps=20, save_summary_steps=20)

est = tf.estimator.DNNLinearCombinedClassifier(
    linear_feature_columns=wide_columns,
    dnn_feature_columns=deep_columns,
    weight_column='weight',
    dnn_hidden_units=hidden_units,
    dnn_optimizer=my_dnn_optimizer,
    dnn_dropout=dropout,
    model_dir=model_dir,
    config=config,
    n_classes=4)

train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn, max_steps=2000)
eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn)
tf.estimator.train_and_evaluate(est, train_spec, eval_spec)

feature_spec = tf.feature_column.make_parse_example_spec(feature_columns)
serving_input_receiver_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(feature_spec)
export_dir = est.export_saved_model(model_dir, serving_input_receiver_fn)

## 加载
est = tf.estimator.DNNLinearCombinedClassifier(
    linear_feature_columns=wide_columns,
    dnn_feature_columns=deep_columns,
    weight_column='weight',
    dnn_hidden_units=hidden_units,
    dnn_optimizer=my_dnn_optimizer,
    dnn_dropout=dropout,
    model_dir=model_dir,
    n_classes=4)
## 预测
result = est.predict(eval_input_fn)
output = []
for prediction in zip(result):
    output.append(int(prediction[0]['classes']))
## 保存  
serving_input_receiver_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(feature_spec)
export_dir = est.export_saved_model(model_dir, serving_input_receiver_fn)

imported = tf.saved_model.load(export_dir)
f = imported.signatures["predict"]
def predict(dataframe):
    for row in dataframe.itertuples():
        feature_map = {}
        feature_map[l] = tf.train.Feature(float_list=tf.train.FloatList(value=[getattr(row,l)]))
        example = tf.train.Example(
            features=tf.train.Features(
                feature = feature_map
            )
        )
        ex = tf.constant([example.SerializeToString()])
        result = f(examples=ex)
    return result