step1. 创建一个网络
model_fn = model_fn_builder(
bert_config=bert_config,
num_labels=len(label_list),
init_checkpoint=FLAGS.init_checkpoint,
learning_rate=FLAGS.learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps,
use_tpu=FLAGS.use_tpu,
use_one_hot_embeddings=FLAGS.use_tpu)
step2.把这个忘了传给estimator,并设置好训练的参数。虽然这里是TPUEstimator,但是原理是一样的
estimator = tf.contrib.tpu.TPUEstimator(
use_tpu=FLAGS.use_tpu,
model_fn=model_fn,
config=run_config,
train_batch_size=FLAGS.train_batch_size,
eval_batch_size=FLAGS.eval_batch_size,
predict_batch_size=FLAGS.predict_batch_size)
step3,开始训练。train_input_fn是训练数据集(整个数据集,不用batch)。
estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
result = estimator.evaluate(input_fn=eval_input_fn, steps=eval_steps)
result = estimator.predict(input_fn=predict_input_fn)