# coding=utf-8
# Copyright 2019 The Interval Bound Propagation Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Helper to keep track of the different losses."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import sonnet as snt
import tensorflow.compat.v1 as tf
# Used to pick the least violated specification.
_BIG_NUMBER = 1e25
ScalarMetrics = collections.namedtuple('ScalarMetrics', [
'nominal_accuracy',
'verified_accuracy',
'attack_accuracy',
'attack_success'])
ScalarLosses = collections.namedtuple('ScalarLosses', [
'nominal_cross_entropy',
'attack_cross_entropy',
'verified_loss'])
class Losses(snt.AbstractModule):
"""Helper to compute our losses."""
def __init__(self, predictor, specification=None, pgd_attack=None,
interval_bounds_loss_type='xent',
interval_bounds_hinge_margin=10.,
label_smoothing=0.):
super(Losses, self).__init__(name='losses')
self._predictor = predictor
self._specification = specification
self._attack = pgd_attack
# Loss type can be any combination of:
# xent: cross-entropy loss
# hinge: hinge loss
# softplus: softplus loss
# with
# all: using all specifications.
# most: using only the specification that is the most violated.
# least: using only the specification that is the least violated.
# random_n: using a random subset of the specifications.
# E.g.: "xent_max" or "hinge_random_3".
tokens = interval_bounds_loss_type.split('_', 1)
if len(tokens) == 1:
loss_type, loss_mode = tokens[0], 'all'
else:
loss_type, loss_mode = tokens
if loss_mode.startswith('random'):
loss_mode, num_samples = loss_mode.split('_', 1)
self._interval_bounds_loss_n = int(num_samples)
if loss_type not in ('xent', 'hinge', 'softplus'):
raise ValueError('interval_bounds_loss_type must be either "xent", '
'"hinge" or "softplus".')
if loss_mode not in ('all', 'most', 'random', 'least'):
raise ValueError('interval_bounds_loss_type must be followed by either '
'"all", "most", "random_N" or "least".')
self._interval_bounds_loss_type = loss_type
self._interval_bounds_loss_mode = loss_mode
self._interval_bounds_hinge_margin = interval_bounds_hinge_margin
self._label_smoothing = label_smoothing
def _build(self, labels):
self._build_nominal_loss(labels)
self._build_verified_loss(labels)
self._build_attack_loss(labels)
def _build_nominal_loss(self, labels):
"""Build natural cross-entropy loss on clean data."""
# Cross-entropy.
nominal_logits = self._predictor.logits
if self._label_smoothing > 0:
num_classes = nominal_logits.shape[1].value
one_hot_labels = tf.one_hot(labels, num_classes)
smooth_positives = 1. - self._label_smoothing
smooth_negatives = self._label_smoothing / num_classes
one_hot_labels = one_hot_labels * smooth_positives + smooth_negatives
nominal_cross_entropy = tf.nn.softmax_cross_entropy_with_logits_v2(
labels=one_hot_labels, logits=nominal_logits)
self._one_hot_labels = one_hot_labels
else:
nominal_cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=labels, logits=nominal_logits)
self._cross_entropy = tf.reduce_mean(nominal_cross_entropy)
# Accuracy.
nominal_correct_examples = tf.equal(labels, tf.argmax(nominal_logits, 1))
self._nominal_accuracy = tf.reduce_mean(
tf.cast(nominal_correct_examples, tf.float32))
def _get_specification_bounds(self):
"""Get upper bounds on specification. Used for building verified loss."""
ibp_bounds = self._specification(self._predictor.modules)
# Compute verified accuracy using IBP bounds.
v = tf.reduce_max(ibp_bounds, axis=1)
self._interval_bounds_accuracy = tf.reduce_mean(
tf.cast(v <= 0., tf.float32))
return ibp_bounds
def _build_verified_loss(self, labels):
"""Build verified loss using an upper bound on specification."""
if not self._specification:
self._verified_loss = tf.constant(0.)
self._interval_bounds_accuracy = tf.constant(0.)
return
# Interval bounds.
bounds = self._get_specification_bounds()
# Select specifications.
if self._interval_bounds_loss_mode == 'all':
pass # Keep bounds the way it is.
elif self._interval_bounds_loss_mode == 'most':
bounds = tf.reduce_max(bounds, axis=1, keepdims=True)
elif self._interval_bounds_loss_mode == 'random':
idx = tf.random.uniform(
[tf.shape(bounds)[0], self._interval_bounds_loss_n],
0, tf.shape(bounds)[1], dtype=tf.int32)
bounds = tf.batch_gather(bounds, idx)
else:
assert self._interval_bounds_loss_mode == 'least'
# This picks the least violated contraint.
mask = tf.cast(bounds < 0., tf.float32)
smallest_violation = tf.reduce_min(
bounds + mask * _BIG_NUMBER, axis=1, keepdims=True)
has_violations = tf.less(
tf.reduce_sum(mask, axis=1, keepdims=True) + .5,
tf.cast(tf.shape(bounds)[1], tf.float32))
largest_bounds = tf.reduce_max(bounds, axis=1, keepdims=True)
bounds = tf.where(has_violations, smallest_violation, largest_bounds)
if self._interval_bounds_loss_type == 'xent':
v = tf.concat(
[bounds, tf.zeros([tf.shape(bounds)[0], 1], dtype=bounds.dtype)],
axis=1)
l = tf.concat(
[tf.zeros_like(bounds),
tf.ones([tf.shape(bounds)[0], 1], dtype=bounds.dtype)],
axis=1)
self._verified_loss = tf.reduce_mean(
tf.nn.softmax_cross_entropy_with_logits_v2(
labels=tf.stop_gradient(l), logits=v))
elif self._interval_bounds_loss_type == 'softplus':
self._verified_loss = tf.reduce_mean(
tf.nn.softplus(bounds + self._interval_bounds_hinge_margin))
else:
assert self._interval_bounds_loss_type == 'hinge'
self._verified_loss = tf.reduce_mean(
tf.maximum(bounds, -self._interval_bounds_hinge_margin))
def _build_attack_loss(self, labels):
"""Build adversarial loss using PGD attack."""
# PGD attack.
if not self._attack:
self._attack_accuracy = tf.constant(0.)
self._attack_success = tf.constant(1.)
self._attack_cross_entropy = tf.constant(0.)
return
if not isinstance(self._predictor.inputs, tf.Tensor):
raise ValueError('Multiple inputs is not supported.')
self._attack(self._predictor.inputs, labels)
correct_examples = tf.equal(labels, tf.argmax(self._attack.logits, 1))
self._attack_accuracy = tf.reduce_mean(
tf.cast(correct_examples, tf.float32))
self._attack_success = tf.reduce_mean(
tf.cast(self._attack.success, tf.float32))
if self._label_smoothing > 0:
attack_cross_entropy = tf.nn.softmax_cross_entropy_with_logits_v2(
labels=self._one_hot_labels, logits=self._attack.logits)
else:
attack_cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=labels, logits=self._attack.logits)
self._attack_cross_entropy = tf.reduce_mean(attack_cross_entropy)
@property
def scalar_metrics(self):
self._ensure_is_connected()
return ScalarMetrics(self._nominal_accuracy,
self._interval_bounds_accuracy,
self._attack_accuracy,
self._attack_success)
@property
def scalar_losses(self):
self._ensure_is_connected()
return ScalarLosses(self._cross_entropy,
self._attack_cross_entropy,
self._verified_loss)
