<commit_before>"""Tests for morph_net.tools.configurable_ops."""

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

import json
import os

from absl import flags

from absl.testing import parameterized

from morph_net.tools import configurable_ops as ops

import tensorflow as tf

from tensorflow.contrib.slim.nets import resnet_utils
from tensorflow.contrib.slim.nets import resnet_v1
from tensorflow.contrib.slim.nets import resnet_v2

resnet_utils = tf.contrib.slim.nets.resnet_utils
resnet_v1 = tf.contrib.slim.nets.resnet_v1
resnet_v2 = tf.contrib.slim.nets.resnet_v2

arg_scope = tf.contrib.framework.arg_scope
add_arg_scope = tf.contrib.framework.add_arg_scope

FLAGS = flags.FLAGS


@add_arg_scope
def mock_fully_connected(*args, **kwargs):
  return {'mock_name': 'myfully_connected', 'args': args, 'kwargs': kwargs}


@add_arg_scope
def mock_conv2d(*args, **kwargs):
  return {'mock_name': 'myconv2d', 'args': args, 'kwargs': kwargs}


@add_arg_scope
def mock_separable_conv2d(*args, **kwargs):
  return {'mock_name': 'myseparable_conv2d', 'args': args, 'kwargs': kwargs}


@add_arg_scope
def mock_concat(*args, **kwargs):
  return {'mock_name': 'myconcat', 'args': args, 'kwargs': kwargs}


@add_arg_scope
def mock_add_n(*args, **kwargs):
  return {'mock_name': 'myaddn', 'args': args, 'kwargs': kwargs}


class ConfigurableOpsTest(parameterized.TestCase, tf.test.TestCase):

  def setUp(self):
    super(ConfigurableOpsTest, self).setUp()
    tf.reset_default_graph()
    self.inputs_shape = [2, 4, 4, 3]
    self.inputs = tf.ones(self.inputs_shape, dtype=tf.float32)
    self.fc_inputs = tf.ones([3, 12])

  def testMapBinding(self):
    # TODO(e1): Clean up this file/test. Split to different tests
    function_dict = {
        'fully_connected': mock_fully_connected,
        'conv2d': mock_conv2d,
        'separable_conv2d': mock_separable_conv2d,
        'concat': mock_concat,
        'add_n': mock_add_n,
    }
    parameterization = {
        'fc/MatMul': 13,
        'conv/Conv2D': 15,
        'sep/separable_conv2d': 17
    }
    num_outputs = lambda res: res['args'][1]
    decorator = ops.ConfigurableOps(
        parameterization=parameterization, function_dict=function_dict)
    fc = decorator.fully_connected(self.fc_inputs, num_outputs=88, scope='fc')
    self.assertEqual('myfully_connected', fc['mock_name'])
    self.assertEqual(parameterization['fc/MatMul'], num_outputs(fc))

    conv2d = decorator.conv2d(
        self.inputs, num_outputs=11, kernel_size=3, scope='conv')
    self.assertEqual('myconv2d', conv2d['mock_name'])
    self.assertEqual(parameterization['conv/Conv2D'], num_outputs(conv2d))

    separable_conv2d = decorator.separable_conv2d(
        self.inputs, num_outputs=88, kernel_size=3, scope='sep')
    self.assertEqual('myseparable_conv2d', separable_conv2d['mock_name'])
    self.assertEqual(parameterization['sep/separable_conv2d'],
                     num_outputs(separable_conv2d))

    concat = decorator.concat(axis=1, values=[1, None, 2])
    self.assertEqual(concat['args'][0], [1, 2])
    self.assertEqual(concat['kwargs']['axis'], 1)
    with self.assertRaises(ValueError):
      _ = decorator.concat(inputs=[1, None, 2])

    add_n = decorator.add_n(name='add_n', inputs=[1, None, 2])
    self.assertEqual(add_n['args'][0], [1, 2])

  def testFullyConnectedOpAllKwargs(self):
    decorator = ops.ConfigurableOps(parameterization={'test/MatMul': 13})
    output = decorator.fully_connected(
        inputs=self.fc_inputs, num_outputs=88, scope='test')
    self.assertEqual(13, output.shape.as_list()[-1])

  def testFullyConnectedOpInputArgs(self):
    decorator = ops.ConfigurableOps(parameterization={'test/MatMul': 14})
    output = decorator.fully_connected(
        self.fc_inputs, num_outputs=87, scope='test')
    self.assertEqual(14, output.shape.as_list()[-1])

  def testFullyConnectedOpAllArgs(self):
    decorator = ops.ConfigurableOps(parameterization={'test/MatMul': 15})
    output = decorator.fully_connected(self.fc_inputs, 86, scope='test')
    self.assertEqual(15, output.shape.as_list()[-1])

  def testSeparableConv2dOp(self):
    parameterization = {'test/separable_conv2d': 12}
    decorator = ops.ConfigurableOps(parameterization=parameterization)
    output = decorator.separable_conv2d(
        self.inputs,
        num_outputs=88,
        kernel_size=3,
        depth_multiplier=1.0,
        scope='test')
    self.assertEqual(12, output.shape.as_list()[-1])

  def testComplexNet(self):
    parameterization = {'Branch0/Conv_1x1/Conv2D': 13, 'Conv3_1x1/Conv2D': 77}
    decorator = ops.ConfigurableOps(parameterization=parameterization)

    def conv2d(inputs, num_outputs, kernel_size, scope):
      return decorator.conv2d(
          inputs, num_outputs=num_outputs, kernel_size=kernel_size, scope=scope)

    net = self.inputs

    with tf.variable_scope('Branch0'):
      branch_0 = conv2d(net, 1, 1, scope='Conv_1x1')
    with tf.variable_scope('Branch1'):
      branch_1 = conv2d(net, 2, 1, scope='Conv_1x1')
      out_2 = conv2d(branch_1, 3, 3, scope='Conv_3x3')
    net = conv2d(net, 1, 1, scope='Conv3_1x1')
    output = tf.concat([net, branch_0, branch_1, out_2], -1)
    expected_output_shape = self.inputs_shape
    expected_output_shape[-1] = 95
    self.assertEqual(expected_output_shape, output.shape.as_list())
    self.assertEqual(2, decorator.constructed_ops['Branch1/Conv_1x1/Conv2D'])<commit_msg>Fixed cases where tf.TensorShape was constructed with float dimensions<commit_after>"""Tests for morph_net.tools.configurable_ops."""

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

import json
import os

from absl import flags

from absl.testing import parameterized

from morph_net.tools import configurable_ops as ops

import tensorflow as tf

from tensorflow.contrib.slim.nets import resnet_utils
from tensorflow.contrib.slim.nets import resnet_v1
from tensorflow.contrib.slim.nets import resnet_v2

resnet_utils = tf.contrib.slim.nets.resnet_utils
resnet_v1 = tf.contrib.slim.nets.resnet_v1
resnet_v2 = tf.contrib.slim.nets.resnet_v2

arg_scope = tf.contrib.framework.arg_scope
add_arg_scope = tf.contrib.framework.add_arg_scope

FLAGS = flags.FLAGS


@add_arg_scope
def mock_fully_connected(*args, **kwargs):
  return {'mock_name': 'myfully_connected', 'args': args, 'kwargs': kwargs}


@add_arg_scope
def mock_conv2d(*args, **kwargs):
  return {'mock_name': 'myconv2d', 'args': args, 'kwargs': kwargs}


@add_arg_scope
def mock_separable_conv2d(*args, **kwargs):
  return {'mock_name': 'myseparable_conv2d', 'args': args, 'kwargs': kwargs}


@add_arg_scope
def mock_concat(*args, **kwargs):
  return {'mock_name': 'myconcat', 'args': args, 'kwargs': kwargs}


@add_arg_scope
def mock_add_n(*args, **kwargs):
  return {'mock_name': 'myaddn', 'args': args, 'kwargs': kwargs}


class ConfigurableOpsTest(parameterized.TestCase, tf.test.TestCase):

  def setUp(self):
    super(ConfigurableOpsTest, self).setUp()
    tf.reset_default_graph()
    self.inputs_shape = [2, 4, 4, 3]
    self.inputs = tf.ones(self.inputs_shape, dtype=tf.float32)
    self.fc_inputs = tf.ones([3, 12])

  def testMapBinding(self):
    # TODO(e1): Clean up this file/test. Split to different tests
    function_dict = {
        'fully_connected': mock_fully_connected,
        'conv2d': mock_conv2d,
        'separable_conv2d': mock_separable_conv2d,
        'concat': mock_concat,
        'add_n': mock_add_n,
    }
    parameterization = {
        'fc/MatMul': 13,
        'conv/Conv2D': 15,
        'sep/separable_conv2d': 17
    }
    num_outputs = lambda res: res['args'][1]
    decorator = ops.ConfigurableOps(
        parameterization=parameterization, function_dict=function_dict)
    fc = decorator.fully_connected(self.fc_inputs, num_outputs=88, scope='fc')
    self.assertEqual('myfully_connected', fc['mock_name'])
    self.assertEqual(parameterization['fc/MatMul'], num_outputs(fc))

    conv2d = decorator.conv2d(
        self.inputs, num_outputs=11, kernel_size=3, scope='conv')
    self.assertEqual('myconv2d', conv2d['mock_name'])
    self.assertEqual(parameterization['conv/Conv2D'], num_outputs(conv2d))

    separable_conv2d = decorator.separable_conv2d(
        self.inputs, num_outputs=88, kernel_size=3, scope='sep')
    self.assertEqual('myseparable_conv2d', separable_conv2d['mock_name'])
    self.assertEqual(parameterization['sep/separable_conv2d'],
                     num_outputs(separable_conv2d))

    concat = decorator.concat(axis=1, values=[1, None, 2])
    self.assertEqual(concat['args'][0], [1, 2])
    self.assertEqual(concat['kwargs']['axis'], 1)
    with self.assertRaises(ValueError):
      _ = decorator.concat(inputs=[1, None, 2])

    add_n = decorator.add_n(name='add_n', inputs=[1, None, 2])
    self.assertEqual(add_n['args'][0], [1, 2])

  def testFullyConnectedOpAllKwargs(self):
    decorator = ops.ConfigurableOps(parameterization={'test/MatMul': 13})
    output = decorator.fully_connected(
        inputs=self.fc_inputs, num_outputs=88, scope='test')
    self.assertEqual(13, output.shape.as_list()[-1])

  def testFullyConnectedOpInputArgs(self):
    decorator = ops.ConfigurableOps(parameterization={'test/MatMul': 14})
    output = decorator.fully_connected(
        self.fc_inputs, num_outputs=87, scope='test')
    self.assertEqual(14, output.shape.as_list()[-1])

  def testFullyConnectedOpAllArgs(self):
    decorator = ops.ConfigurableOps(parameterization={'test/MatMul': 15})
    output = decorator.fully_connected(self.fc_inputs, 86, scope='test')
    self.assertEqual(15, output.shape.as_list()[-1])

  def testSeparableConv2dOp(self):
    parameterization = {'test/separable_conv2d': 12}
    decorator = ops.ConfigurableOps(parameterization=parameterization)
    output = decorator.separable_conv2d(
        self.inputs,
        num_outputs=88,
        kernel_size=3,
        depth_multiplier=1,
        scope='test')
    self.assertEqual(12, output.shape.as_list()[-1])

  def testComplexNet(self):
    parameterization = {'Branch0/Conv_1x1/Conv2D': 13, 'Conv3_1x1/Conv2D': 77}
    decorator = ops.ConfigurableOps(parameterization=parameterization)

    def conv2d(inputs, num_outputs, kernel_size, scope):
      return decorator.conv2d(
          inputs, num_outputs=num_outputs, kernel_size=kernel_size, scope=scope)

    net = self.inputs

    with tf.variable_scope('Branch0'):
      branch_0 = conv2d(net, 1, 1, scope='Conv_1x1')
    with tf.variable_scope('Branch1'):
      branch_1 = conv2d(net, 2, 1, scope='Conv_1x1')
      out_2 = conv2d(branch_1, 3, 3, scope='Conv_3x3')
    net = conv2d(net, 1, 1, scope='Conv3_1x1')
    output = tf.concat([net, branch_0, branch_1, out_2], -1)
    expected_output_shape = self.inputs_shape
    expected_output_shape[-1] = 95
    self.assertEqual(expected_output_shape, output.shape.as_list())
    self.assertEqual(2, decorator.constructed_ops['Branch1/Conv_1x1/Conv2D'])
{'fraction_non_alphanumeric': 0.13039919978175865, 'fraction_numerical': 0.02791670455578794, 'mean_word_length': 4.614088820826953, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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}}
{'dir': 'git-commits-cleaned', 'id': '3734760', 'n_tokens_mistral': 3872, 'n_tokens_neox': 3737, 'n_words': 675}