| 516 | return feature_cols |
| 517 | |
| 518 | def train(sess_config, |
| 519 | input_hooks, |
| 520 | model, |
| 521 | data_init_op, |
| 522 | steps, |
| 523 | checkpoint_dir, |
| 524 | tf_config=None, |
| 525 | server=None): |
| 526 | model.is_training = True |
| 527 | hooks = [] |
| 528 | hooks.extend(input_hooks) |
| 529 | |
| 530 | scaffold = tf.train.Scaffold( |
| 531 | local_init_op=tf.group(tf.local_variables_initializer(), data_init_op), |
| 532 | saver=tf.train.Saver(max_to_keep=args.keep_checkpoint_max, sharded=True)) |
| 533 | |
| 534 | stop_hook = tf.train.StopAtStepHook(last_step=steps) |
| 535 | log_hook = tf.train.LoggingTensorHook( |
| 536 | { |
| 537 | 'steps': model.global_step, |
| 538 | 'loss': model.loss |
| 539 | }, every_n_iter=100) |
| 540 | hooks.append(stop_hook) |
| 541 | hooks.append(log_hook) |
| 542 | if args.timeline > 0: |
| 543 | hooks.append( |
| 544 | tf.train.ProfilerHook(save_steps=args.timeline, |
| 545 | output_dir=checkpoint_dir)) |
| 546 | save_steps = args.save_steps if args.save_steps or args.no_eval else steps |
| 547 | ''' |
| 548 | Incremental_Checkpoint |
| 549 | Please add `save_incremental_checkpoint_secs` in 'tf.train.MonitoredTrainingSession' |
| 550 | it's default to None, Incremental_save checkpoint time in seconds can be set |
| 551 | to use incremental checkpoint function, like `tf.train.MonitoredTrainingSession( |
| 552 | save_incremental_checkpoint_secs=args.incremental_ckpt)` |
| 553 | ''' |
| 554 | if args.incremental_ckpt and not args.tf: |
| 555 | print("Incremental_Checkpoint is not really enabled.") |
| 556 | print("Please see the comments in the code.") |
| 557 | sys.exit() |
| 558 | |
| 559 | with tf.train.MonitoredTrainingSession( |
| 560 | master=server.target if server else '', |
| 561 | is_chief=tf_config['is_chief'] if tf_config else True, |
| 562 | hooks=hooks, |
| 563 | scaffold=scaffold, |
| 564 | checkpoint_dir=checkpoint_dir, |
| 565 | save_checkpoint_steps=save_steps, |
| 566 | summary_dir=checkpoint_dir, |
| 567 | save_summaries_steps=args.save_steps, |
| 568 | config=sess_config) as sess: |
| 569 | while not sess.should_stop(): |
| 570 | sess.run([model.loss, model.train_op]) |
| 571 | print("Training completed.") |
| 572 | |
| 573 | def eval(sess_config, input_hooks, model, data_init_op, steps, checkpoint_dir): |
| 574 | model.is_training = False |