Loss defined by string works out-of-the-box but not tf.keras.losses.X

[samedii]

The following works: loss='binary_crossentropy', but stops working when replaced by: tf.keras.losses.BinaryCrossentropy()

Are we expected to define weights? Or are we maybe supposed to wrap our loss functions in nsl.keras.adversarial_loss?

Edit: Can define weight key with sample_weight_key=
Edit2: I realized that it's actually the shape of values that is incorrect. This seems like a bug then?

Traceback (most recent call last):
  File "test_nsl3.py", line 86, in <module>
    main()
  File "test_nsl3.py", line 78, in main
    adv_model, train_dataset, test_dataset)
  File "test_nsl3.py", line 67, in train_and_evaluate
    model.fit(train_dataset.batch(16), epochs=200, steps_per_epoch=100)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training.py", line 728, in fit
    use_multiprocessing=use_multiprocessing)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 224, in fit
    distribution_strategy=strategy)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 547, in _process_training_inputs
    use_multiprocessing=use_multiprocessing)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 594, in _process_inputs
    steps=steps)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training.py", line 2419, in _standardize_user_data
    all_inputs, y_input, dict_inputs = self._build_model_with_inputs(x, y)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training.py", line 2622, in _build_model_with_inputs
    self._set_inputs(cast_inputs)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/training.py", line 2709, in _set_inputs
    outputs = self(inputs, **kwargs)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 842, in __call__
    outputs = call_fn(cast_inputs, *args, **kwargs)
  File "/home/richard/Documents/cloud/venv/lib/python3.6/site-packages/tensorflow_core/python/autograph/impl/api.py", line 237, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in converted code:
    relative to /home/richard/Documents/cloud/venv/lib/python3.6/site-packages:

    neural_structured_learning/keras/adversarial_regularization.py:623 call  *
        adv_loss = adversarial_loss(
    neural_structured_learning/keras/adversarial_regularization.py:139 adversarial_loss  *
        adv_loss = loss_fn(labels, adv_output, adv_sample_weights)
    neural_structured_learning/keras/adversarial_regularization.py:570 _compute_total_loss  *
        loss, _ = _compute_loss_and_metrics(self._labeled_losses, None, labels,
    neural_structured_learning/keras/adversarial_regularization.py:372 _compute_loss_and_metrics  *
        loss_value = loss(label, output, sample_weights)
    neural_structured_learning/keras/adversarial_regularization.py:175 __call__  *
        loss_value = super(_LossWrapper, self).__call__(*args, **kwargs)
    tensorflow_core/python/keras/losses.py:128 __call__
        losses, sample_weight, reduction=self._get_reduction())
    tensorflow_core/python/keras/utils/losses_utils.py:107 compute_weighted_loss
        losses, sample_weight)
    tensorflow_core/python/ops/losses/util.py:148 scale_losses_by_sample_weight
        sample_weight = weights_broadcast_ops.broadcast_weights(sample_weight, losses)
    tensorflow_core/python/ops/weights_broadcast_ops.py:167 broadcast_weights
        with ops.control_dependencies((assert_broadcastable(weights, values),)):
    tensorflow_core/python/ops/weights_broadcast_ops.py:103 assert_broadcastable
        weights_rank_static, values.shape, weights.shape))

    ValueError: weights can not be broadcast to values. values.rank=0. weights.rank=2. values.shape=(). weights.shape=(None, 1).

Here is code reproducing the error:

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import neural_structured_learning as nsl
import tensorflow as tf


def prepare_datasets():
  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

  def make_dataset(x, y, shuffle=False):
    y = y.reshape((-1, 28, 28, 1)).astype('float32') / 255.0
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    if shuffle:
      dataset = dataset.repeat().shuffle(100)
    return dataset.map(to_dict)

  return make_dataset(y_train, x_train, True), make_dataset(y_test, x_test)


def to_dict(x, y):
    return {'feature': x, 'label': y}


def build_base_model():
  return tf.keras.Sequential([
    tf.keras.Input(shape=(1,), dtype=tf.float32, name='feature'),
    tf.keras.layers.Dense(28*28, activation='relu'),
    tf.keras.layers.Dense(28*28, activation='sigmoid'),
    tf.keras.layers.Lambda(lambda x:  tf.reshape(x, (-1, 28, 28, 1))),
  ])


def apply_adversarial_regularization(model):
  adv_config = nsl.configs.make_adv_reg_config(
      multiplier=0.5,
      adv_step_size=0.5,
      adv_grad_norm='infinity'
  )
  return nsl.keras.AdversarialRegularization(
      model, label_keys=['label'], adv_config=adv_config)


def build_adv_model():
  base_model = build_base_model()
  return apply_adversarial_regularization(base_model)


def train_and_evaluate(model, train_dataset, test_dataset):
  model.compile(
      optimizer='adam',
      loss=tf.keras.losses.BinaryCrossentropy(),
      # loss='binary_crossentropy',
      metrics=['accuracy']
  )
  model.fit(train_dataset.batch(16), epochs=200, steps_per_epoch=100)
  eval_result = model.evaluate(test_dataset)
  return list(zip(model.metrics_names, eval_result))


def main():

  train_dataset, test_dataset = prepare_datasets()

  adv_model = build_adv_model()
  adv_result = train_and_evaluate(
      adv_model, train_dataset, test_dataset)

  for metric_name, result in adv_result:
    print('Eval %s for adversarial model: %s' % (metric_name, result))



if __name__ == '__main__':
  main()

[samedii]

Related to #8? Will try to install master instead. Edit: Master does not seem to solve this.

[csferng]

@samedii, Thanks for reporting the issue and the code for reproduction.

Yes, this is a bug, which is caused by AdversarialRegularization that creates a wrapper around the given Loss object (in this case BinaryCrossentropy); the wrapper expects to receive unaggregated loss values from the Loss object but actually gets an aggregated value.

We are working on fixing this issue. Meanwhile, a workaround is to set the reduction of the Loss object to NONE:

adv_model.compile(
    # ...
    loss=tf.keras.losses.BinaryCrossentropy(reduction=tf.losses.Reduction.NONE),
    # ...
)

Note that the feature tensor (x) has to be in float type so that AdversarialRegularization can be effective. You might want to add x = x.astype('float32') into your make_dataset function.

[samedii]

Thanks for the quick answer! Ah I see why this is needed for adversarial learning. Imo I think it's fine to just give an easier error message for new people and document this requirement

[csferng]

The fix is checked in and will be included in the next release. The workaround will no longer be needed afterward.