text: # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
  #
  # 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.
  # ==============================================================================
  """Tests for locally-connected layers."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  
  from absl.testing import parameterized
  
  from tensorflow.python import keras
  from tensorflow.python.framework import test_util as tf_test_util
  from tensorflow.python.keras import testing_utils
  from tensorflow.python.platform import test
  from tensorflow.python.training.rmsprop import RMSPropOptimizer
  
  
  _DATA_FORMAT_PADDING_IMPLEMENTATION = [{
      'data_format': 'channels_first',
      'padding': 'valid',
      'implementation': 1
  }, {
      'data_format': 'channels_first',
      'padding': 'same',
      'implementation': 1
  }, {
      'data_format': 'channels_last',
      'padding': 'valid',
      'implementation': 1
  }, {
      'data_format': 'channels_last',
      'padding': 'same',
      'implementation': 1
  }, {
      'data_format': 'channels_first',
      'padding': 'valid',
      'implementation': 2
  }, {
      'data_format': 'channels_first',
      'padding': 'same',
      'implementation': 2
  }, {
      'data_format': 'channels_last',
      'padding': 'valid',
      'implementation': 2
  }, {
      'data_format': 'channels_last',
      'padding': 'same',
      'implementation': 2
  }, {
      'data_format': 'channels_first',
      'padding': 'valid',
      'implementation': 3
  }, {
      'data_format': 'channels_first',
      'padding': 'same',
      'implementation': 3
  }, {
      'data_format': 'channels_last',
      'padding': 'valid',
      'implementation': 3
  }, {
      'data_format': 'channels_last',
      'padding': 'same',
      'implementation': 3
  }]
  
  
  @tf_test_util.run_all_in_graph_and_eager_modes
  class LocallyConnected1DLayersTest(test.TestCase, parameterized.TestCase):
  
    @parameterized.parameters(_DATA_FORMAT_PADDING_IMPLEMENTATION)
    def test_locallyconnected_1d(self, data_format, padding, implementation):
      with self.cached_session():
        num_samples = 2
        num_steps = 8
        input_dim = 5
        filter_length = 3
        filters = 4
  
        for strides in [1]:
          if padding == 'same' and strides != 1:
            continue
          kwargs = {
              'filters': filters,
              'kernel_size': filter_length,
              'padding': padding,
              'strides': strides,
              'data_format': data_format,
              'implementation': implementation
          }
  
          if padding == 'same' and implementation == 1:
            self.assertRaises(ValueError, keras.layers.LocallyConnected1D,
                              **kwargs)
          else:
            testing_utils.layer_test(
                keras.layers.LocallyConnected1D,
                kwargs=kwargs,
                input_shape=(num_samples, num_steps, input_dim))
  
    @parameterized.parameters(_DATA_FORMAT_PADDING_IMPLEMENTATION)
    def test_locallyconnected_1d_regularization(self, data_format, padding,
                                                implementation):
      num_samples = 2
      num_steps = 8
      input_dim = 5
      filter_length = 3
      filters = 4
      kwargs = {
          'filters': filters,
          'kernel_size': filter_length,
          'kernel_regularizer': 'l2',
          'bias_regularizer': 'l2',
          'activity_regularizer': 'l2',
          'data_format': data_format,
          'implementation': implementation,
          'padding': padding
      }
  
      if padding == 'same' and implementation == 1:
        self.assertRaises(ValueError, keras.layers.LocallyConnected1D, **kwargs)
      else:
        with self.cached_session():
          layer = keras.layers.LocallyConnected1D(**kwargs)
          layer.build((num_samples, num_steps, input_dim))
          self.assertEqual(len(layer.losses), 2)
          layer(
              keras.backend.variable(
                  np.ones((num_samples, num_steps, input_dim))))
          self.assertEqual(len(layer.losses), 3)
  
        k_constraint = keras.constraints.max_norm(0.01)
        b_constraint = keras.constraints.max_norm(0.01)
        kwargs = {
            'filters': filters,
            'kernel_size': filter_length,
            'kernel_constraint': k_constraint,
            'bias_constraint': b_constraint,
        }
        with self.cached_session():
          layer = keras.layers.LocallyConnected1D(**kwargs)
          layer.build((num_samples, num_steps, input_dim))
          self.assertEqual(layer.kernel.constraint, k_constraint)
          self.assertEqual(layer.bias.constraint, b_constraint)
  
  
  @tf_test_util.run_all_in_graph_and_eager_modes
  class LocallyConnected2DLayersTest(test.TestCase, parameterized.TestCase):
  
    @parameterized.parameters(_DATA_FORMAT_PADDING_IMPLEMENTATION)
    def test_locallyconnected_2d(self, data_format, padding, implementation):
      with self.cached_session():
        num_samples = 8
        filters = 3
        stack_size = 4
        num_row = 6
        num_col = 10
  
        for strides in [(1, 1), (2, 2)]:
          if padding == 'same' and strides != (1, 1):
            continue
  
          kwargs = {
              'filters': filters,
              'kernel_size': 3,
              'padding': padding,
              'kernel_regularizer': 'l2',
              'bias_regularizer': 'l2',
              'strides': strides,
              'data_format': data_format,
              'implementation': implementation
          }
  
          if padding == 'same' and implementation == 1:
            self.assertRaises(ValueError, keras.layers.LocallyConnected2D,
                              **kwargs)
          else:
            testing_utils.layer_test(
                keras.layers.LocallyConnected2D,
                kwargs=kwargs,
                input_shape=(num_samples, num_row, num_col, stack_size))
  
    @parameterized.parameters(_DATA_FORMAT_PADDING_IMPLEMENTATION)
    def test_locallyconnected_2d_channels_first(self, data_format, padding,
                                                implementation):
      with self.cached_session():
        num_samples = 8
        filters = 3
        stack_size = 4
        num_row = 6
        num_col = 10
        kwargs = {
            'filters': filters,
            'kernel_size': 3,
            'data_format': data_format,
            'implementation': implementation,
            'padding': padding
        }
  
        if padding == 'same' and implementation == 1:
          self.assertRaises(ValueError, keras.layers.LocallyConnected2D, **kwargs)
        else:
          testing_utils.layer_test(
              keras.layers.LocallyConnected2D,
              kwargs=kwargs,
              input_shape=(num_samples, num_row, num_col, stack_size))
  
    @parameterized.parameters(_DATA_FORMAT_PADDING_IMPLEMENTATION)
    def test_locallyconnected_2d_regularization(self, data_format, padding,
                                                implementation):
      num_samples = 2
      filters = 3
      stack_size = 4
      num_row = 6
      num_col = 7
      kwargs = {
          'filters': filters,
          'kernel_size': 3,
          'kernel_regularizer': 'l2',
          'bias_regularizer': 'l2',
          'activity_regularizer': 'l2',
          'implementation': implementation,
          'padding': padding,
          'data_format': data_format
      }
  
      if padding == 'same' and implementation == 1:
        self.assertRaises(ValueError, keras.layers.LocallyConnected2D, **kwargs)
      else:
        with self.cached_session():
          layer = keras.layers.LocallyConnected2D(**kwargs)
          layer.build((num_samples, num_row, num_col, stack_size))
          self.assertEqual(len(layer.losses), 2)
          layer(
              keras.backend.variable(
                  np.ones((num_samples, num_row, num_col, stack_size))))
          self.assertEqual(len(layer.losses), 3)
  
        k_constraint = keras.constraints.max_norm(0.01)
        b_constraint = keras.constraints.max_norm(0.01)
        kwargs = {
            'filters': filters,
            'kernel_size': 3,
            'kernel_constraint': k_constraint,
            'bias_constraint': b_constraint,
        }
        with self.cached_session():
          layer = keras.layers.LocallyConnected2D(**kwargs)
          layer.build((num_samples, num_row, num_col, stack_size))
          self.assertEqual(layer.kernel.constraint, k_constraint)
          self.assertEqual(layer.bias.constraint, b_constraint)
  
  
  @tf_test_util.run_all_in_graph_and_eager_modes
  class LocallyConnectedImplementationModeTest(test.TestCase,
                                               parameterized.TestCase):
  
    @parameterized.parameters([
        {'width': 1, 'data_format': 'channels_first'},
        {'width': 1, 'data_format': 'channels_last'},
        {'width': 6, 'data_format': 'channels_first'},
        {'width': 6, 'data_format': 'channels_last'},
    ])
    def test_locallyconnected_implementation(self, width, data_format):
      with self.cached_session():
        num_samples = 4
        num_classes = 3
        num_epochs = 2
  
        np.random.seed(1)
        targets = np.random.randint(0, num_classes, (num_samples,))
  
        height = 7
        filters = 2
        inputs = get_inputs(data_format, filters, height, num_samples, width)
  
        kernel_x = (3,)
        kernel_y = () if width == 1 else (2,)
        stride_x = (1,)
        stride_y = () if width == 1 else (3,)
        layers = 2
  
        kwargs = {
            'layers': layers,
            'filters': filters,
            'kernel_size': kernel_x + kernel_y,
            'strides': stride_x + stride_y,
            'data_format': data_format,
            'num_classes': num_classes
        }
  
        model_1 = get_model(implementation=1, **kwargs)
        model_2 = get_model(implementation=2, **kwargs)
        model_3 = get_model(implementation=3, **kwargs)
  
        # Build models.
        model_1.train_on_batch(inputs, targets)
        model_2.train_on_batch(inputs, targets)
        model_3.train_on_batch(inputs, targets)
  
        # Copy weights.
        copy_model_weights(model_from=model_2, model_to=model_1)
        copy_model_weights(model_from=model_2, model_to=model_3)
  
        # Compare outputs at initialization.
        out_1 = model_1.call(inputs)
        out_2 = model_2.call(inputs)
        out_3 = model_3.call(inputs)
  
        self.assertAllCloseAccordingToType(
            out_2, out_1, rtol=1e-5, atol=1e-5)
        self.assertAllCloseAccordingToType(
            out_2, out_3, rtol=1e-5, atol=1e-5)
        self.assertAllCloseAccordingToType(
            out_1, out_3, rtol=1e-5, atol=1e-5)
  
        # Train.
        model_1.fit(
            x=inputs,
            y=targets,
            epochs=num_epochs,
            batch_size=num_samples)
        model_2.fit(
            x=inputs,
            y=targets,
            epochs=num_epochs,
            batch_size=num_samples)
        model_3.fit(
            x=inputs,
            y=targets,
            epochs=num_epochs,
            batch_size=num_samples)
  
        # Compare outputs after a few training steps.
        out_1 = model_1.call(inputs)
        out_2 = model_2.call(inputs)
        out_3 = model_3.call(inputs)
  
        self.assertAllCloseAccordingToType(
            out_2, out_1, atol=2e-4)
        self.assertAllCloseAccordingToType(
            out_2, out_3, atol=2e-4)
        self.assertAllCloseAccordingToType(
            out_1, out_3, atol=2e-4)
  
    def test_make_2d(self):
      input_shapes = [
          (0,),
          (0, 0),
          (1,),
          (2,),
          (3,),
          (1, 0),
          (0, 3),
          (1, 1),
          (1, 2),
          (3, 1),
          (2, 2),
          (3, 3),
          (1, 0, 1),
          (5, 2, 3),
          (3, 5, 6, 7, 0),
          (3, 2, 2, 4, 4),
          (1, 2, 3, 4, 7, 2),
      ]
      np.random.seed(1)
  
      for input_shape in input_shapes:
        inputs = np.random.normal(0, 1, input_shape)
        inputs_tf = keras.backend.variable(inputs)
  
        split_dim = np.random.randint(0, inputs.ndim + 1)
        shape_2d = (int(np.prod(inputs.shape[:split_dim])),
                    int(np.prod(inputs.shape[split_dim:])))
        inputs_2d = np.reshape(inputs, shape_2d)
  
        inputs_2d_tf = keras.layers.local.make_2d(inputs_tf, split_dim)
        inputs_2d_tf = keras.backend.get_value(inputs_2d_tf)
  
        self.assertAllCloseAccordingToType(inputs_2d, inputs_2d_tf)
  
  
  def get_inputs(data_format, filters, height, num_samples, width):
    if data_format == 'channels_first':
      if width == 1:
        input_shape = (filters, height)
      else:
        input_shape = (filters, height, width)
  
    elif data_format == 'channels_last':
      if width == 1:
        input_shape = (height, filters)
      else:
        input_shape = (height, width, filters)
  
    else:
      raise NotImplementedError(data_format)
  
    inputs = np.random.normal(0, 1,
                              (num_samples,) + input_shape).astype(np.float32)
    return inputs
  
  
  def xent(y_true, y_pred):
    y_true = keras.backend.cast(
        keras.backend.reshape(y_true, (-1,)),
        keras.backend.dtypes_module.int32)
  
    return keras.backend.nn.sparse_softmax_cross_entropy_with_logits(
        labels=y_true,
        logits=y_pred)
  
  
  def get_model(implementation,
                filters,
                kernel_size,
                strides,
                layers,
                num_classes,
                data_format):
    model = keras.Sequential()
  
    if len(kernel_size) == 1:
      lc_layer = keras.layers.LocallyConnected1D
    elif len(kernel_size) == 2:
      lc_layer = keras.layers.LocallyConnected2D
    else:
      raise NotImplementedError(kernel_size)
  
    for _ in range(layers):
      model.add(lc_layer(
          padding='valid',
          kernel_initializer=keras.initializers.random_normal(),
          bias_initializer=keras.initializers.random_normal(),
          filters=filters,
          strides=strides,
          kernel_size=kernel_size,
          activation=keras.activations.relu,
          data_format=data_format,
          implementation=implementation))
  
    model.add(keras.layers.Flatten())
    model.add(keras.layers.Dense(num_classes))
    model.compile(
        optimizer=RMSPropOptimizer(0.01),
        metrics=[keras.metrics.categorical_accuracy],
        loss=xent
    )
    return model
  
  
  def copy_lc_weights_2_to_1(lc_layer_2_from, lc_layer_1_to):
    lc_2_kernel, lc_2_bias = lc_layer_2_from.weights
    lc_2_kernel_masked = lc_2_kernel * lc_layer_2_from.kernel_mask
  
    data_format = lc_layer_2_from.data_format
  
    if data_format == 'channels_first':
      if isinstance(lc_layer_2_from, keras.layers.LocallyConnected1D):
        permutation = (3, 0, 1, 2)
      elif isinstance(lc_layer_2_from, keras.layers.LocallyConnected2D):
        permutation = (4, 5, 0, 1, 2, 3)
      else:
        raise NotImplementedError(lc_layer_2_from)
  
    elif data_format == 'channels_last':
      if isinstance(lc_layer_2_from, keras.layers.LocallyConnected1D):
        permutation = (2, 0, 1, 3)
      elif isinstance(lc_layer_2_from, keras.layers.LocallyConnected2D):
        permutation = (3, 4, 0, 1, 2, 5)
      else:
        raise NotImplementedError(lc_layer_2_from)
  
    else:
      raise NotImplementedError(data_format)
  
    lc_2_kernel_masked = keras.backend.permute_dimensions(
        lc_2_kernel_masked, permutation)
  
    lc_2_kernel_mask = keras.backend.math_ops.not_equal(
        lc_2_kernel_masked, 0)
    lc_2_kernel_flat = keras.backend.array_ops.boolean_mask(
        lc_2_kernel_masked, lc_2_kernel_mask)
    lc_2_kernel_reshaped = keras.backend.reshape(lc_2_kernel_flat,
                                                 lc_layer_1_to.kernel.shape)
  
    lc_2_kernel_reshaped = keras.backend.get_value(lc_2_kernel_reshaped)
    lc_2_bias = keras.backend.get_value(lc_2_bias)
  
    lc_layer_1_to.set_weights([lc_2_kernel_reshaped, lc_2_bias])
  
  
  def copy_lc_weights_2_to_3(lc_layer_2_from, lc_layer_3_to):
    lc_2_kernel, lc_2_bias = lc_layer_2_from.weights
    lc_2_kernel_masked = lc_2_kernel * lc_layer_2_from.kernel_mask
  
    lc_2_kernel_masked = keras.layers.local.make_2d(
        lc_2_kernel_masked, split_dim=keras.backend.ndim(lc_2_kernel_masked) // 2)
    lc_2_kernel_masked = keras.backend.transpose(lc_2_kernel_masked)
    lc_2_kernel_mask = keras.backend.math_ops.not_equal(lc_2_kernel_masked, 0)
    lc_2_kernel_flat = keras.backend.array_ops.boolean_mask(
        lc_2_kernel_masked, lc_2_kernel_mask)
  
    lc_2_kernel_flat = keras.backend.get_value(lc_2_kernel_flat)
    lc_2_bias = keras.backend.get_value(lc_2_bias)
  
    lc_layer_3_to.set_weights([lc_2_kernel_flat, lc_2_bias])
  
  
  def copy_model_weights(model_from, model_to):
    for l in range(len(model_from.layers)):
      layer_from = model_from.layers[l]
      layer_to = model_to.layers[l]
  
      if (isinstance(
          layer_from,
          (keras.layers.LocallyConnected2D, keras.layers.LocallyConnected1D)) and
          isinstance(layer_to, (keras.layers.LocallyConnected2D,
                                keras.layers.LocallyConnected1D))):
        if layer_from.implementation == 2:
          if layer_to.implementation == 1:
            copy_lc_weights_2_to_1(layer_from, layer_to)
          elif layer_to.implementation == 3:
            copy_lc_weights_2_to_3(layer_from, layer_to)
          else:
            raise NotImplementedError
  
        else:
          raise NotImplementedError
  
      elif isinstance(layer_from, keras.layers.Dense):
        weights_2, bias_2 = layer_from.weights
        weights_2 = keras.backend.get_value(weights_2)
        bias_2 = keras.backend.get_value(bias_2)
        layer_to.set_weights([weights_2, bias_2])
  
      else:
        continue
  
  
  if __name__ == '__main__':
    test.main()
  