<jupyter_start><jupyter_text>[View in Colaboratory](https://colab.research.google.com/github/ZacCranko/robustlearningexperiments/blob/master/regulariser_unit_tests.ipynb)<jupyter_code># Copyright 2018 Google LLC and <NAME>

# 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

#     https://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.

import numpy as np
import tensorflow as tf
# Copyright 2018 Google LLC.
# SPDX-License-Identifier: Apache-2.0
def power_iterate_conv(layer, num_iter):
  """Perform power iteration for a convolutional layer."""
  assert isinstance(layer, tf.keras.layers.Conv2D)
  weights = layer.kernel
  strides = (1,) + layer.strides + (1,)
  padding = layer.padding.upper()
  
  with tf.variable_scope(None, default_name='power_iteration'):
    u_var = tf.get_variable(
       'u_conv', [1] + map(int, layer.output_shape[1:]),
       initializer=tf.random_normal_initializer(),
       trainable=False)
    u = u_var
    
    for _ in xrange(num_iter):
      v = tf.nn.conv2d_transpose(
         u, weights, [1] + map(int, layer.input_shape[1:]), strides, padding)
      v /= tf.sqrt(tf.maximum(2 * tf.nn.l2_loss(v), 1e-12))
      u = tf.nn.conv2d(v, weights, strides, padding)
      u /= tf.sqrt(tf.maximum(2 * tf.nn.l2_loss(u), 1e-12))
      
    tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, tf.assign(u_var, u))

    u = tf.stop_gradient(u)
    v = tf.stop_gradient(v)
    return tf.reduce_sum(u * tf.nn.conv2d(v, weights, strides, padding))
  
def power_iterate_dense(layer, num_iter):
  """Perform power iteration for a fully connected layer."""
  assert isinstance(layer, tf.keras.layers.Dense)
  weights = layer.kernel
  output_shape, input_shape = weights.get_shape().as_list()

  with tf.variable_scope(None, default_name='power_iteration'):
    u_var = tf.get_variable(
       'u',  map(int, [output_shape]) + [1],
       initializer=tf.random_normal_initializer(),
       trainable=False)
    u = u_var

    for _ in xrange(num_iter):
      v = tf.matmul(weights, u, transpose_a=True)
      v /= tf.sqrt(tf.maximum(2 * tf.nn.l2_loss(v), 1e-12))
      u = tf.matmul(weights, v)
      u /= tf.sqrt(tf.maximum(2 * tf.nn.l2_loss(u), 1e-12))

    tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, tf.assign(u_var, u))

    u = tf.stop_gradient(u)
    v = tf.stop_gradient(v)
    return tf.reduce_sum(u * tf.matmul(weights, v))
def operator_norm(layer, ord = 2, **kwargs):
    """Compute operator norm for a Keras layer."""
    
    with tf.variable_scope(None, default_name='operator_norm'):
      if ord == 1:
          w = layer.kernel
          if isinstance(layer, tf.keras.layers.Conv2D):
              sum_w = tf.reduce_sum(tf.abs(w), [0, 1, 3])
          else:
              sum_w = tf.reduce_sum(tf.abs(w), 1)

          return tf.reduce_max(sum_w)

      elif ord == 2:
          num_iter = kwargs.get('num_iter', 5)
          if isinstance(layer, tf.keras.layers.Conv2D):
              return power_iterate_conv(layer, num_iter)
          else:
              return power_iterate_dense(layer, num_iter)

      elif ord == np.inf:
          w = layer.kernel
          if isinstance(layer, tf.keras.layers.Conv2D):
              sum_w = tf.reduce_sum(tf.abs(w), [0, 1, 2])
          else:
              sum_w = tf.reduce_sum(tf.abs(w), 0)

          return tf.reduce_max(sum_w)
def tf_assert_almost_equal(actual, desired, **kwargs):
  with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    _actual  = actual.eval()
    _desired = desired.eval()

  return np.testing.assert_almost_equal(_actual, _desired, **kwargs)

def conv_matrix(layer):
  """Build the matrix associated with the convolution operation."""
  assert isinstance(layer, tf.keras.layers.Conv2D)
  with tf.variable_scope(None, default_name='build_conv_matrix'):
    weights = layer.kernel
    strides = (1,) + layer.strides + (1,)
    padding = layer.padding.upper()
    in_h,  in_w,  in_ch  = layer.input_shape[1:4]
    out_h, out_w, out_ch = layer.output_shape[1:4]

    id_mx     = tf.reshape(tf.eye(in_h*in_w*in_ch), 
                            (in_h*in_w*in_ch, in_h, in_w, in_ch))
    conv_mx_t = tf.reshape(tf.nn.conv2d(id_mx, weights, strides, padding), 
                            (in_h*in_w*in_ch, out_h*out_w*out_ch))
    return tf.transpose(conv_mx_t)
model = tf.keras.Sequential()
conv1 = tf.keras.layers.Conv2D(32, 5, 1, padding='SAME',
                               input_shape=(28, 28, 1))
model.add(conv1)
model.add(tf.keras.layers.Activation('relu'))
model.add(tf.keras.layers.MaxPooling2D(2, 2, padding='SAME'))

conv2 = tf.keras.layers.Conv2D(64, 5, 1, padding='SAME')
model.add(conv2)

dense1 = tf.keras.layers.Dense(1024)
model.add(dense1)
model.add(tf.keras.layers.Activation('relu'))

dense2 = tf.keras.layers.Dense(10)
model.add(dense2)
num_iter = 200

# test dense layers
for layer in ['dense1', 'dense2']:
  _layer = eval(layer)
  
  op = "1_norm"
  # axis = 1 here since _layer.kernel is stored transposed for dense layers
  inf_opn_mx = tf.reduce_max(tf.reduce_sum(tf.abs(_layer.kernel), axis = 1))
  tf_assert_almost_equal(operator_norm(_layer, 1), inf_opn_mx, err_msg = "%s(%s)"%(op,layer), decimal = 1)
  
  op = "inf_norm"
  # axis = 0 here since _layer.kernel is stored transposed for dense layers
  inf_opn_mx = tf.reduce_max(tf.reduce_sum(tf.abs(_layer.kernel), axis = 0))
  tf_assert_almost_equal(operator_norm(_layer, np.inf), inf_opn_mx, err_msg = "%s(%s)"%(op,layer), decimal = 1)
  
  op = "spectral_norm"
  spec_pow  = operator_norm(_layer, 2, num_iter = num_iter)
  spec_svd  = tf.svd(_layer.kernel, compute_uv=False)
  tf_assert_almost_equal(spec_pow, spec_svd[0], decimal = 2, err_msg = "%s(%s)"%(op,layer))

# test conv layers
for layer in ['conv1', 'conv2']:
  _layer = eval(layer)
  
  op = "1_norm"
  conv_mx = conv_matrix(_layer)
  desired = tf.reduce_max(tf.reduce_sum(tf.abs(conv_mx), axis = 0))
  tf_assert_almost_equal(operator_norm(_layer, 1), desired, err_msg = "%s(%s)"%(op,layer), decimal = 1)
  
  op = "inf_norm"
  conv_mx = conv_matrix(_layer)
  desired = tf.reduce_max(tf.reduce_sum(tf.abs(conv_mx), axis = 1))
  tf_assert_almost_equal(operator_norm(_layer, np.inf), desired, err_msg = "%s(%s)"%(op,layer), decimal = 1)
  
  op = "spectral_norm"
  spec_pow  = operator_norm(_layer, 2, num_iter = num_iter)
  spec_svd  = tf.svd(conv_matrix(_layer), compute_uv=False)
  tf_assert_almost_equal(spec_pow, spec_svd[0], err_msg = "%s(%s)"%(op,layer), decimal = 2)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13718673040706938, 'fraction_numerical': 0.01839779805881501, 'mean_word_length': 3.566137566137566, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 6, 'https://': 2, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '495926', 'license': 'Apache-2.0', 'n_tokens_mistral': 2623, 'n_tokens_neox': 2500, 'n_words': 527, 'path': 'regulariser_unit_tests.ipynb', 'repo_name': 'ZacCranko/robustlearningexperiments', 'size': '6907'}