nsl.estimator.add_graph_regularization

Adds graph regularization to a tf.estimator.Estimator.

nsl.estimator.add_graph_regularization(
    estimator, embedding_fn, optimizer_fn=None, graph_reg_config=None
)

Used in the tutorials
Graph-based Neural Structured Learning in TFX

Args
`estimator`	An object of type `tf.estimator.Estimator`.
`embedding_fn`	A function that accepts the input layer (dictionary of feature names and corresponding batched tensor values) as its first argument and an instance of `tf.estimator.ModeKeys` as its second argument to indicate if the mode is training, evaluation, or prediction, and returns the corresponding embeddings or logits to be used for graph regularization.
`optimizer_fn`	A function that accepts no arguments and returns an instance of `tf.train.Optimizer`.
`graph_reg_config`	An instance of `nsl.configs.GraphRegConfig` that specifies various hyperparameters for graph regularization.

Returns
A modified `tf.estimator.Estimator` object with graph regularization incorporated into its loss.

Args