integrated_gradients_test.py

# coding=utf-8
# Copyright 2020 The Google Research 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.

# Lint as: python2, python3
"""Tests for google_research.attribution.integrated_gradients."""

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

import tensorflow.compat.v1 as tf

from attribution import integrated_gradients
from tensorflow.contrib import layers as contrib_layers


class AttributionTest(tf.test.TestCase):

  def testAddIntegratedGradientsOps(self):
    with tf.Graph().as_default() as graph:
      var1 = tf.compat.v1.get_variable(
          name='var1', initializer=[[[1., 2., 3.]]])
      input_tensor = tf.placeholder(shape=[None, None, 3], dtype=tf.float32)
      x = tf.multiply(input_tensor, [[[1.]]])
      var1_times_x = tf.multiply(var1, x)
      var2 = tf.compat.v1.get_variable(
          name='var2', initializer=[[4., 5.], [6., 7.], [4., 3.]])
      matmul = tf.einsum('ijk,kl->ijl', var1_times_x, var2)
      output_tensor = tf.reduce_sum(matmul, [1, 2], name='out')
      input_feed_dict = {input_tensor.name: [[[2., 3., 4.], [5., 6., 7.]]]}
      num_evals = tf.placeholder_with_default(
          tf.constant(20, name='num_evals'), shape=())
      attribution_hooks = integrated_gradients.AddIntegratedGradientsOps(
          graph=graph,
          attribution_tensors=[x],
          output_tensor=output_tensor,
          num_evals=num_evals,
          attribution_dims_map={x: [1]},
          zero_baseline_tensors=set([x]))

      with tf.Session(graph=graph) as sess:
        sess.run(tf.global_variables_initializer())
        result = sess.run(
            attribution_hooks['mean_grads'],
            attribution_hooks['create_combined_feed_dict'](input_feed_dict))
        self.assertTupleEqual(result[0].shape, (2, 1))
        self.assertAlmostEqual(result[0][0, 0], 180.)
        self.assertAlmostEqual(result[0][1, 0], 348.)

  def testAddBOWIntegratedGradientsOps(self):
    with tf.Graph().as_default() as graph:
      # pyformat: disable
      embedding_weights = tf.constant([[1., 3., 5.],
                                       [4., 6., 8.],
                                       [4., 5., 4.]])
      batch_size = tf.placeholder_with_default(tf.constant(1, tf.int64), [])
      sparse_ids = tf.SparseTensor(
          [[0, 0], [0, 1], [0, 2]], [2, 0, 2], [batch_size, 3])
      # pyformat: enable
      sparse_embedding = contrib_layers.safe_embedding_lookup_sparse(
          embedding_weights,
          sparse_ids,
          combiner='sum',
          partition_strategy='div')

      vector = tf.constant([1., 2., 4.])
      output_tensor = tf.reduce_sum(vector * sparse_embedding, axis=1)
      embedding_lookup = tf.nn.embedding_lookup(
          embedding_weights,
          tf.sparse_tensor_to_dense(sparse_ids),
          partition_strategy='div')
      bow_attribution_hooks = integrated_gradients.AddBOWIntegratedGradientsOps(
          graph, [embedding_lookup], [sparse_embedding], [], output_tensor)
      with tf.Session(graph=graph) as sess:
        sess.run(tf.global_variables_initializer())
        result = sess.run(bow_attribution_hooks['bow_attributions'])
        self.assertTupleEqual(result[0].shape, (3,))
        # Since the output is a sum of dot products, attributions are simply dot
        # products of the embedding with [1., 2., 4.].
        self.assertAlmostEqual(result[0][0], 30.)
        self.assertAlmostEqual(result[0][1], 27.)
        self.assertAlmostEqual(result[0][2], 30.)

  def testGetEmbeddingLookupList(self):
    with tf.Graph().as_default() as graph:
      # pyformat: disable
      embedding_weights = tf.constant([[1., 2., 3.],
                                       [4., 6., 8.],
                                       [5., 4., 3.]])
      batch_size = tf.placeholder_with_default(tf.constant(1, tf.int64), [])
      sparse_ids1 = tf.SparseTensor(
          [[0, 0], [0, 1]], [2, 0], [batch_size, 2])
      sparse_weights1 = tf.SparseTensor(
          [[0, 0], [0, 1]], [4., 3.], [batch_size, 2])
      sparse_ids2 = tf.SparseTensor(
          [[0, 0], [0, 1], [0, 2]], [1, 1, 2], [batch_size, 3])
      # pyformat: enable
      embedding_lookup_list = integrated_gradients.GetEmbeddingLookupList(
          ['feature1', 'feature2'],
          {'feature1': embedding_weights, 'feature2': embedding_weights},
          {'feature1': sparse_ids1, 'feature2': sparse_ids2},
          sparse_weights={'feature1': sparse_weights1, 'feature2': None},
          combiners={'feature1': 'sqrtn', 'feature2': 'mean'})
      with tf.Session(graph=graph) as sess:
        sess.run(tf.global_variables_initializer())
        result = sess.run(embedding_lookup_list)
        expected = [[[[4., 3.2, 2.4],
                      [0.6, 1.2, 1.8]]],
                    [[[1.33, 2., 2.66],
                      [1.33, 2., 2.66],
                      [1.66, 1.33, 1.]]]]
        self.assertTupleEqual(result[0].shape, (1, 2, 3))
        self.assertTupleEqual(result[1].shape, (1, 3, 3))
        for idx in range(len(expected)):
          for row in range(len(expected[idx][0])):
            for col in range(len(expected[idx][0][row])):
              self.assertAlmostEqual(result[idx][0, row, col],
                                     expected[idx][0][row][col],
                                     delta=0.1)


if __name__ == '__main__':
  tf.test.main()