# ======================== 构建输入 ========================
# 配置tfrecord的数据结构格式
name2features = {}
for f in sparse_features:
name2features[f] = tf.io.FixedLenFeature([], tf.int64)
for f in dense_features:
name2features[f] = tf.io.FixedLenFeature([], tf.float32)
for f in [args.target]:
name2features[f] = tf.io.FixedLenFeature([], tf.float32)
if args.mode == 'train':
train_input_fn = tfrecord2fn(os.path.join(args.tfrecord_dir,
'train.tfrecord'),
name2features,
args.batch_size,
args.num_epoches,
drop_remainder=True,
mode=tf.estimator.ModeKeys.TRAIN,
target=args.target)
elif args.mode == 'eval':
eval_input_fn = tfrecord2fn(os.path.join(args.tfrecord_dir,
'eval.tfrecord'),
name2features,
args.batch_size,
args.num_epoches,
drop_remainder=True,
mode=tf.estimator.ModeKeys.EVAL,
target=args.target)
elif args.mode == 'train_eval':
train_input_fn = tfrecord2fn(os.path.join(args.tfrecord_dir,
'train.tfrecord'),
"cp": tf.FixedLenFeature([], tf.int64),
"trestbps": tf.FixedLenFeature([], tf.int64),
"chol": tf.FixedLenFeature([], tf.int64),
"fbs": tf.FixedLenFeature([], tf.int64),
"restecg": tf.FixedLenFeature([], tf.int64),
"thalach": tf.FixedLenFeature([], tf.int64),
"exang": tf.FixedLenFeature([], tf.int64),
"oldpeak": tf.FixedLenFeature([], tf.int64),
"slope": tf.FixedLenFeature([], tf.int64),
"ca": tf.FixedLenFeature([], tf.int64),
"thal": tf.FixedLenFeature([], tf.int64),
"target": tf.FixedLenFeature([], tf.float32),
}
if split:
train_input_fn = tfrecord2fn('heart.tfrecord.train', name2features, batch_size, num_epochs,drop_remainder=True, is_training=True, target='target')
eval_input_fn = tfrecord2fn('heart.tfrecord.eval', name2features, batch_size, num_epochs, drop_remainder=True, is_training=True, target='target')
else:
train_input_fn = tfrecord2fn('heart.tfrecord', name2features, batch_size, num_epochs,drop_remainder=True, is_training=True, target='target')
eval_input_fn = tfrecord2fn('heart.tfrecord', name2features, batch_size, num_epochs, drop_remainder=True, is_training=True, target='target')
test_input_fn = tfrecord2fn('heart.tfrecord', name2features, batch_size, num_epochs, drop_remainder=True, is_training=False)
# ======================== 进行训练 ========================
early_stopping_hook = tf.estimator.experimental.stop_if_no_decrease_hook(
estimator=estimator,
metric_name='eval_loss',
max_steps_without_decrease=1000,
min_steps=10,
run_every_secs=None,
run_every_steps=1000