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