def _model_fn(self): """Returns a model_fn.""" def _train_op_fn(loss): """Defines train op used in ranking head.""" update_ops = tf.compat.v1.get_collection( tf.compat.v1.GraphKeys.UPDATE_OPS) minimize_op = self._optimizer.minimize( loss=loss, global_step=tf.compat.v1.train.get_global_step()) train_op = tf.group([update_ops, minimize_op]) return train_op ranking_head = head.create_ranking_head( loss_fn=losses.make_loss_fn( self._hparams.get("loss"), weights_feature_name=self._hparams.get(_LOSS_WEIGHT), reduction=self._loss_reduction), eval_metric_fns=self._eval_metric_fns(), train_op_fn=_train_op_fn) return model.make_groupwise_ranking_fn( group_score_fn=self._group_score_fn, group_size=1, transform_fn=self._transform_fn, ranking_head=ranking_head)
def setUp(self): super(GroupwiseRankingEstimatorTest, self).setUp() ops.reset_default_graph() self._model_dir = test.get_temp_dir() gfile.MakeDirs(self._model_dir) model_fn = model.make_groupwise_ranking_fn( _group_score_fn, group_size=2, transform_fn=feature.make_identity_transform_fn(['context', 'weight']), ranking_head=head.create_ranking_head( loss_fn=losses.make_loss_fn( losses.RankingLossKey.PAIRWISE_HINGE_LOSS, weights_feature_name='weight'), optimizer=training.AdagradOptimizer(learning_rate=0.1))) self._estimator = estimator.Estimator(model_fn, self._model_dir)