The `model_fn` for TPUEstimator.
(features, labels, mode, params)
| 560 | """Returns `model_fn` closure for TPUEstimator.""" |
| 561 | |
| 562 | def model_fn(features, labels, mode, params): # pylint: disable=unused-argument |
| 563 | """The `model_fn` for TPUEstimator.""" |
| 564 | |
| 565 | tf.logging.info("*** Features ***") |
| 566 | for name in sorted(features.keys()): |
| 567 | tf.logging.info(" name = %s, shape = %s" % (name, features[name].shape)) |
| 568 | |
| 569 | input_ids = features["input_ids"] |
| 570 | |
| 571 | is_training = (mode == tf.estimator.ModeKeys.TRAIN) |
| 572 | |
| 573 | model = GroverModel( |
| 574 | config=config, |
| 575 | is_training=is_training, |
| 576 | input_ids=input_ids, |
| 577 | pad_token_id=config.pad_token_id, |
| 578 | chop_off_last_token=True, |
| 579 | ) |
| 580 | |
| 581 | total_loss = model.lm_loss() |
| 582 | print(model.logits_flat) |
| 583 | print(total_loss) |
| 584 | |
| 585 | if is_training: |
| 586 | train_op, train_metrics = create_optimizer( |
| 587 | total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu) |
| 588 | tvars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) |
| 589 | else: |
| 590 | train_op = None |
| 591 | train_metrics = {} |
| 592 | tvars = tf.trainable_variables() |
| 593 | params_sum = np.sum([np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()]) |
| 594 | tf.logging.info("**** Trainable params_sum ****") |
| 595 | tf.logging.info(params_sum) |
| 596 | initialized_variable_names = {} |
| 597 | scaffold_fn = None |
| 598 | if init_checkpoint: |
| 599 | (assignment_map, initialized_variable_names |
| 600 | ) = get_assignment_map_from_checkpoint(tvars, init_checkpoint) |
| 601 | if use_tpu: |
| 602 | def tpu_scaffold(): |
| 603 | tf.train.init_from_checkpoint(init_checkpoint, assignment_map) |
| 604 | return tf.train.Scaffold() |
| 605 | |
| 606 | scaffold_fn = tpu_scaffold |
| 607 | else: |
| 608 | tf.train.init_from_checkpoint(init_checkpoint, assignment_map) |
| 609 | |
| 610 | tf.logging.info("**** Trainable Variables ****") |
| 611 | for var in tvars: |
| 612 | init_string = "" |
| 613 | if var.name in initialized_variable_names: |
| 614 | init_string = ", *INIT_FROM_CKPT*" |
| 615 | tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, |
| 616 | init_string) |
| 617 | |
| 618 | output_spec = None |
| 619 | if mode == tf.estimator.ModeKeys.TRAIN: |
nothing calls this directly
no test coverage detected