def __init__(self,
base_model,
label_keys=('label', ),
sample_weight_key=None,
adv_config=None):
"""Constructor of `AdversarialRegularization` class.
Args:
base_model: A `tf.Keras.Model` to which adversarial regularization will be
applied.
label_keys: A tuple of strings denoting which keys in the input features
(a `dict` mapping keys to tensors) represent labels. This list should be
1-to-1 corresponding to the output of the `base_model`.
sample_weight_key: A string denoting which key in the input feature (a
`dict` mapping keys to tensors) represents sample weight. If not set,
the weight is 1.0 for each input example.
adv_config: Instance of `nsl.configs.AdvRegConfig` for configuring
adversarial regularization.
"""
super(AdversarialRegularization,
self).__init__(name='AdversarialRegularization')
self.base_model = base_model
self.label_keys = label_keys
self.sample_weight_key = sample_weight_key
self.adv_config = adv_config or nsl_configs.AdvRegConfig()
def __init__(self,
base_model,
label_keys=('label', ),
sample_weight_key=None,
adv_config=None,
base_with_labels_in_features=False):
"""Constructor of `AdversarialRegularization` class.
Args:
base_model: A `tf.Keras.Model` to which adversarial regularization will be
applied.
label_keys: A tuple of strings denoting which keys in the input features
(a `dict` mapping keys to tensors) represent labels. This list should be
1-to-1 corresponding to the output of the `base_model`.
sample_weight_key: A string denoting which key in the input feature (a
`dict` mapping keys to tensors) represents sample weight. If not set,
the weight is 1.0 for each input example.
adv_config: Instance of `nsl.configs.AdvRegConfig` for configuring
adversarial regularization.
base_with_labels_in_features: A Boolean value indicating whether the base
model expects label features as input. This option is effective only
when the base model is a subclassed Keras model. (For functional and
Sequential models, the expected inputs can be inferred from the model
itself.) If set to true, the base model will be called with an input
dictionary including label and sample-weight features. If set to false,
label and sample-weight features will not present in base model's input
dictionary.
"""
super(AdversarialRegularization,
self).__init__(name='AdversarialRegularization')
self.base_model = base_model
self.label_keys = label_keys
self.sample_weight_key = sample_weight_key
self.adv_config = adv_config or nsl_configs.AdvRegConfig()
self._base_with_labels_in_features = base_with_labels_in_features
def adversarial_loss(features,
labels,
model,
loss_fn,
sample_weights=None,
adv_config=None,
predictions=None,
labeled_loss=None,
gradient_tape=None,
model_kwargs=None):
"""Computes the adversarial loss for `model` given `features` and `labels`.
This utility function adds adversarial perturbations to the input `features`,
runs the `model` on the perturbed features for predictions, and returns the
corresponding loss `loss_fn(labels, model(perturbed_features))`. This function
can be used in a Keras subclassed model and a custom training loop. This can
also be used freely as a helper function in eager execution mode.
The adversarial perturbation is based on the gradient of the labeled loss on
the original input features, i.e. `loss_fn(labels, model(features))`.
Therefore, this function needs to compute the model's predictions on the input
features as `model(features)`, and the labeled loss as `loss_fn(labels,
predictions)`. If predictions or labeled loss have already been computed, they
can be passed in via the `predictions` and `labeled_loss` arguments in order
to save computational resources. Note that in eager execution mode,
`gradient_tape` needs to be set accordingly when passing in `predictions` or
`labeled_loss`, so that the gradient can be computed correctly.
Example:
```python
# A linear regression model (for demonstrating the usage only)
model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(2,))])
loss_fn = tf.keras.losses.MeanSquaredError()
optimizer = tf.keras.optimizers.SGD()
# Custom training loop. (The actual training data is omitted for clarity.)
for x, y in train_dataset:
with tf.GradientTape() as tape_w:
# A separate GradientTape is needed for watching the input.
with tf.GradientTape() as tape_x:
tape_x.watch(x)
# Regular forward pass.
labeled_loss = loss_fn(y, model(x))
# Calculates the adversarial loss. This will reuse labeled_loss and will
# consume tape_x.
adv_loss = nsl.keras.adversarial_loss(
x, y, model, loss_fn, labeled_loss=labeled_loss, gradient_tape=tape_x)
# Combines both losses. This could also be a weighted combination.
total_loss = labeled_loss + adv_loss
# Regular backward pass.
gradients = tape_w.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
```
Arguments:
features: Input features, should be a `Tensor` or a collection of `Tensor`
objects. If it is a collection, the first dimension of all `Tensor`
objects inside should be the same (i.e. batch size).
labels: Target labels.
model: A callable that takes `features` as inputs and computes `predictions`
as outputs. An example would be a `tf.keras.Model` object.
loss_fn: A callable which calcualtes labeled loss from `labels`,
`predictions`, and `sample_weights`. An example would be a
`tf.keras.losses.Loss` object.
sample_weights: (optional) A 1-D `Tensor` of weights for the examples, with
the same length as the first dimension of `features`.
adv_config: (optional) An `nsl.configs.AdvRegConfig` object for adversarial
regularization hyperparameters. Use `nsl.configs.make_adv_reg_config` to
construct one.
predictions: (optional) Precomputed value of `model(features)`. If set, the
value will be reused when calculating adversarial regularization. In eager
mode, the `gradient_tape` has to be set as well.
labeled_loss: (optional) Precomputed value of `loss_fn(labels,
model(features))`. If set, the value will be reused when calculating
adversarial regularization. In eager mode, the `gradient_tape` has to be
set as well.
gradient_tape: (optional) A `tf.GradientTape` object watching `features`.
model_kwargs: (optional) A dictionary of additional keyword arguments to be
passed to the `model`.
Returns:
A `Tensor` for adversarial regularization loss, i.e. labeled loss on
adversarially perturbed features.
"""
if adv_config is None:
adv_config = nsl_configs.AdvRegConfig()
if model_kwargs is not None:
model = functools.partial(model, **model_kwargs)
# Calculates labeled_loss if not provided.
if labeled_loss is None:
# Reuses the tape if provided; otherwise creates a new tape.
gradient_tape = gradient_tape or tf.GradientTape()
with gradient_tape:
gradient_tape.watch(tf.nest.flatten(features))
# Calculates prediction if not provided.
predictions = predictions if predictions is not None else model(
features)
labeled_loss = loss_fn(labels, predictions, sample_weights)
adv_input, adv_sample_weights = nsl_lib.gen_adv_neighbor(
features,
labeled_loss,
config=adv_config.adv_neighbor_config,
gradient_tape=gradient_tape)
adv_output = model(adv_input)
if sample_weights is not None:
adv_sample_weights = tf.math.multiply(sample_weights,
adv_sample_weights)
adv_loss = loss_fn(labels, adv_output, adv_sample_weights)
return adv_loss
def add_adversarial_regularization(estimator,
optimizer_fn=None,
adv_config=None):
"""Adds adversarial regularization to a `tf.estimator.Estimator`.
The returned estimator will include the adversarial loss as a regularization
term in its training objective, and will be trained using the optimizer
provided by `optimizer_fn`. `optimizer_fn` (along with the hyperparameters)
should be set to the same one used in the base `estimator`.
If `optimizer_fn` is not set, a default optimizer `tf.train.AdagradOptimizer`
with `learning_rate=0.05` will be used.
Args:
estimator: A `tf.estimator.Estimator` object, the base model.
optimizer_fn: A function that accepts no arguments and returns an instance
of `tf.train.Optimizer`. This optimizer (instead of the one used in
`estimator`) will be used to train the model. If not specified, default to
`tf.train.AdagradOptimizer` with `learning_rate=0.05`.
adv_config: An instance of `nsl.configs.AdvRegConfig` that specifies various
hyperparameters for adversarial regularization.
Returns:
A modified `tf.estimator.Estimator` object with adversarial regularization
incorporated into its loss.
"""
if not adv_config:
adv_config = nsl_configs.AdvRegConfig()
base_model_fn = estimator._model_fn # pylint: disable=protected-access
def adv_model_fn(features, labels, mode, params=None, config=None):
"""The adversarial-regularized model_fn.
Args:
features: This is the first item returned from the `input_fn` passed to
`train`, `evaluate`, and `predict`. This should be a single `tf.Tensor`
or `dict` of same.
labels: This is the second item returned from the `input_fn` passed to
`train`, `evaluate`, and `predict`. This should be a single `tf.Tensor`
or dict of same (for multi-head models). If mode is
`tf.estimator.ModeKeys.PREDICT`, `labels=None` will be passed. If the
`model_fn`'s signature does not accept `mode`, the `model_fn` must still
be able to handle `labels=None`.
mode: Optional. Specifies if this is training, evaluation, or prediction.
See `tf.estimator.ModeKeys`.
params: Optional `dict` of hyperparameters. Will receive what is passed to
Estimator in the `params` parameter. This allows users to configure
Estimators from hyper parameter tuning.
config: Optional `estimator.RunConfig` object. Will receive what is passed
to Estimator as its `config` parameter, or a default value. Allows
setting up things in the model_fn based on configuration such as
`num_ps_replicas`, or `model_dir`. Unused currently.
Returns:
A `tf.estimator.EstimatorSpec` with adversarial regularization.
"""
# Uses the same variable scope for calculating the original objective and
# adversarial regularization.
with tf.compat.v1.variable_scope(tf.compat.v1.get_variable_scope(),
reuse=tf.compat.v1.AUTO_REUSE,
auxiliary_name_scope=False):
# If no 'params' is passed, then it is possible for base_model_fn not to
# accept a 'params' argument. See documentation for tf.estimator.Estimator
# for additional context.
if params:
original_spec = base_model_fn(features, labels, mode, params, config)
else:
original_spec = base_model_fn(features, labels, mode, config)
# Adversarial regularization only happens in training.
if mode != tf.estimator.ModeKeys.TRAIN:
return original_spec
adv_neighbor, _ = nsl_lib.gen_adv_neighbor(features, original_spec.loss,
adv_config.adv_neighbor_config)
# Runs the base model again to compute loss on adv_neighbor.
adv_spec = base_model_fn(adv_neighbor, labels, mode, config)
final_loss = original_spec.loss + adv_config.multiplier * adv_spec.loss
if not optimizer_fn:
# Default to the Adagrad optimizer, the same as canned DNNEstimator.
optimizer = tf.train.AdagradOptimizer(learning_rate=0.05)
else:
optimizer = optimizer_fn()
final_train_op = optimizer.minimize(
loss=final_loss, global_step=tf.compat.v1.train.get_global_step())
return original_spec._replace(loss=final_loss, train_op=final_train_op)
# Replaces the model_fn while keeps other fields/methods in the estimator.
estimator._model_fn = adv_model_fn # pylint: disable=protected-access
return estimator
def add_adversarial_regularization(estimator,
optimizer_fn=None,
adv_config=None):
"""Adds adversarial regularization to a `tf.estimator.Estimator`.
The returned estimator will include the adversarial loss as a regularization
term in its training objective, and will be trained using the optimizer
provided by `optimizer_fn`. `optimizer_fn` (along with the hyperparameters)
should be set to the same one used in the base `estimator`.
If `optimizer_fn` is not set, a default optimizer `tf.train.AdagradOptimizer`
with `learning_rate=0.05` will be used.
Args:
estimator: A `tf.estimator.Estimator` object, the base model.
optimizer_fn: A function that accepts no arguments and returns an instance
of `tf.train.Optimizer`. This optimizer (instead of the one used in
`estimator`) will be used to train the model. If not specified, default to
`tf.train.AdagradOptimizer` with `learning_rate=0.05`.
adv_config: An instance of `nsl.configs.AdvRegConfig` that specifies various
hyperparameters for adversarial regularization.
Returns:
A modified `tf.estimator.Estimator` object with adversarial regularization
incorporated into its loss.
"""
if not adv_config:
adv_config = nsl_configs.AdvRegConfig()
base_model_fn = estimator._model_fn # pylint: disable=protected-access
try:
base_model_fn_args = inspect.signature(base_model_fn).parameters.keys()
except AttributeError: # For Python 2 compatibility
base_model_fn_args = inspect.getargspec(base_model_fn).args # pylint: disable=deprecated-method
def adv_model_fn(features, labels, mode, params=None, config=None):
"""The adversarial-regularized model_fn.
Args:
features: This is the first item returned from the `input_fn` passed to
`train`, `evaluate`, and `predict`. This should be a single `tf.Tensor`
or `dict` of same.
labels: This is the second item returned from the `input_fn` passed to
`train`, `evaluate`, and `predict`. This should be a single `tf.Tensor`
or dict of same (for multi-head models). If mode is
`tf.estimator.ModeKeys.PREDICT`, `labels=None` will be passed. If the
`model_fn`'s signature does not accept `mode`, the `model_fn` must still
be able to handle `labels=None`.
mode: Optional. Specifies if this is training, evaluation, or prediction.
See `tf.estimator.ModeKeys`.
params: Optional `dict` of hyperparameters. Will receive what is passed to
Estimator in the `params` parameter. This allows users to configure
Estimators from hyper parameter tuning.
config: Optional `estimator.RunConfig` object. Will receive what is passed
to Estimator as its `config` parameter, or a default value. Allows
setting up things in the model_fn based on configuration such as
`num_ps_replicas`, or `model_dir`. Unused currently.
Returns:
A `tf.estimator.EstimatorSpec` with adversarial regularization.
"""
# Parameters 'params' and 'config' are optional. If they are not passed,
# then it is possible for base_model_fn not to accept these arguments.
# See documentation for tf.estimator.Estimator for additional context.
kwargs = {'mode': mode}
if 'params' in base_model_fn_args:
kwargs['params'] = params
if 'config' in base_model_fn_args:
kwargs['config'] = config
base_fn = functools.partial(base_model_fn, **kwargs)
# Uses the same variable scope for calculating the original objective and
# adversarial regularization.
with tf.compat.v1.variable_scope(tf.compat.v1.get_variable_scope(),
reuse=tf.compat.v1.AUTO_REUSE,
auxiliary_name_scope=False):
original_spec = base_fn(features, labels)
# Adversarial regularization only happens in training.
if mode != tf.estimator.ModeKeys.TRAIN:
return original_spec
adv_neighbor, _ = nsl_lib.gen_adv_neighbor(
features,
original_spec.loss,
adv_config.adv_neighbor_config,
# The pgd_model_fn is a dummy identity function since loss is
# directly available from base_fn.
pgd_model_fn=lambda features: features,
pgd_loss_fn=lambda labels, features: base_fn(features, labels).
loss,
pgd_labels=labels,
use_while_loop=False)
# Runs the base model again to compute loss on adv_neighbor.
adv_spec = base_fn(adv_neighbor, labels)
scaled_adversarial_loss = adv_config.multiplier * adv_spec.loss
tf.compat.v1.summary.scalar('loss/scaled_adversarial_loss',
scaled_adversarial_loss)
supervised_loss = original_spec.loss
tf.compat.v1.summary.scalar('loss/supervised_loss',
supervised_loss)
final_loss = supervised_loss + scaled_adversarial_loss
if not optimizer_fn:
# Default to the Adagrad optimizer, the same as canned DNNEstimator.
optimizer = tf.compat.v1.train.AdagradOptimizer(
learning_rate=0.05)
else:
optimizer = optimizer_fn()
train_op = optimizer.minimize(
loss=final_loss,
global_step=tf.compat.v1.train.get_global_step())
update_ops = tf.compat.v1.get_collection(
tf.compat.v1.GraphKeys.UPDATE_OPS)
if update_ops:
train_op = tf.group(train_op, *update_ops)
return original_spec._replace(loss=final_loss, train_op=train_op)
# Replaces the model_fn while keeps other fields/methods in the estimator.
estimator._model_fn = adv_model_fn # pylint: disable=protected-access
return estimator
