"""Tests for quantization"""
import numpy as np
import unittest
import os
import shutil
import yaml
import tensorflow as tf


def build_fake_yaml():
    fake_yaml = '''
        model:
          name: fake_yaml
          framework: tensorflow
          inputs: x
          outputs: op_to_store
        device: cpu
        evaluation:
          accuracy:
            metric:
              topk: 1
        tuning:
            strategy:
              name: random
            accuracy_criterion:
              relative: 0.01
            workspace:
              path: saved
        '''
    y = yaml.load(fake_yaml, Loader=yaml.SafeLoader)
    with open('fake_yaml.yaml', "w", encoding="utf-8") as f:
        yaml.dump(y, f)
    f.close()


def build_fake_yaml2():
    fake_yaml = '''
        model:
          name: fake_yaml
          framework: tensorflow
          inputs: x
          outputs: op_to_store
        device: cpu
        evaluation:
          accuracy:
            metric:
              topk: 1
        tuning:
          strategy:
            name: random
          exit_policy:
            max_trials: 5
          accuracy_criterion:
            relative: -0.01
          workspace:
            path: saved
        '''
    y = yaml.load(fake_yaml, Loader=yaml.SafeLoader)
    with open('fake_yaml2.yaml', "w", encoding="utf-8") as f:
        yaml.dump(y, f)
    f.close()


def build_fake_model():
    try:
        graph = tf.Graph()
        graph_def = tf.GraphDef()
        with tf.Session() as sess:
            x = tf.placeholder(tf.float64, shape=(1, 3, 3, 1), name='x')
            y = tf.constant(np.random.random((2, 2, 1, 1)), name='y')
            op = tf.nn.conv2d(input=x, filter=y, strides=[
                              1, 1, 1, 1], padding='VALID', name='op_to_store')

            sess.run(tf.global_variables_initializer())
            constant_graph = tf.graph_util.convert_variables_to_constants(
                sess, sess.graph_def, ['op_to_store'])

        graph_def.ParseFromString(constant_graph.SerializeToString())
        with graph.as_default():
            tf.import_graph_def(graph_def, name='')
    except:
        graph = tf.Graph()
        graph_def = tf.compat.v1.GraphDef()
        with tf.compat.v1.Session() as sess:
            x = tf.compat.v1.placeholder(tf.float64, shape=(1, 3, 3, 1), name='x')
            y = tf.compat.v1.constant(np.random.random((2, 2, 1, 1)), name='y')
            op = tf.nn.conv2d(input=x, filters=y, strides=[
                              1, 1, 1, 1], padding='VALID', name='op_to_store')

            sess.run(tf.compat.v1.global_variables_initializer())
            constant_graph = tf.compat.v1.graph_util.convert_variables_to_constants(sess, sess.graph_def, [
                                                                                    'op_to_store'])

        graph_def.ParseFromString(constant_graph.SerializeToString())
        with graph.as_default():
            tf.import_graph_def(graph_def, name='')
    return graph


class TestQuantization(unittest.TestCase):

    @classmethod
    def setUpClass(self):
        self.constant_graph = build_fake_model()
        build_fake_yaml()
        build_fake_yaml2()

    @classmethod
    def tearDownClass(self):
        os.remove('fake_yaml.yaml')
        os.remove('fake_yaml2.yaml')

        shutil.rmtree("saved", ignore_errors=True)

    def test_ru_random_one_trial(self):
        from neural_compressor.experimental import Quantization, common

        quantizer = Quantization('fake_yaml.yaml')
        dataset = quantizer.dataset('dummy', (100, 3, 3, 1), label=True)
        quantizer.calib_dataloader = common.DataLoader(dataset)
        quantizer.eval_dataloader = common.DataLoader(dataset)
        quantizer.model = self.constant_graph
        quantizer()

    def test_ru_random_max_trials(self):
        from neural_compressor.experimental import Quantization, common

        quantizer = Quantization('fake_yaml2.yaml')
        dataset = quantizer.dataset('dummy', (100, 3, 3, 1), label=True)
        quantizer.calib_dataloader = common.DataLoader(dataset)
        quantizer.eval_dataloader = common.DataLoader(dataset)
        quantizer.model = self.constant_graph
        quantizer()


if __name__ == "__main__":
    unittest.main()

[((619, 663), 'yaml.load', 'yaml.load', (['fake_yaml'], {'Loader': 'yaml.SafeLoader'}), '(fake_yaml, Loader=yaml.SafeLoader)\n', (628, 663), False, 'import yaml\n'), ((1296, 1340), 'yaml.load', 'yaml.load', (['fake_yaml'], {'Loader': 'yaml.SafeLoader'}), '(fake_yaml, Loader=yaml.SafeLoader)\n', (1305, 1340), False, 'import yaml\n'), ((4415, 4430), 'unittest.main', 'unittest.main', ([], {}), '()\n', (4428, 4430), False, 'import unittest\n'), ((735, 750), 'yaml.dump', 'yaml.dump', (['y', 'f'], {}), '(y, f)\n', (744, 750), False, 'import yaml\n'), ((1413, 1428), 'yaml.dump', 'yaml.dump', (['y', 'f'], {}), '(y, f)\n', (1422, 1428), False, 'import yaml\n'), ((1500, 1510), 'tensorflow.Graph', 'tf.Graph', ([], {}), '()\n', (1508, 1510), True, 'import tensorflow as tf\n'), ((1532, 1545), 'tensorflow.GraphDef', 'tf.GraphDef', ([], {}), '()\n', (1543, 1545), True, 'import tensorflow as tf\n'), ((3374, 3401), 'os.remove', 'os.remove', (['"""fake_yaml.yaml"""'], {}), "('fake_yaml.yaml')\n", (3383, 3401), False, 'import os\n'), ((3411, 3439), 'os.remove', 'os.remove', (['"""fake_yaml2.yaml"""'], {}), "('fake_yaml2.yaml')\n", (3420, 3439), False, 'import os\n'), ((3451, 3493), 'shutil.rmtree', 'shutil.rmtree', (['"""saved"""'], {'ignore_errors': '(True)'}), "('saved', ignore_errors=True)\n", (3464, 3493), False, 'import shutil\n'), ((3633, 3663), 'neural_compressor.experimental.Quantization', 'Quantization', (['"""fake_yaml.yaml"""'], {}), "('fake_yaml.yaml')\n", (3645, 3663), False, 'from neural_compressor.experimental import Quantization, common\n'), ((3776, 3802), 'neural_compressor.experimental.common.DataLoader', 'common.DataLoader', (['dataset'], {}), '(dataset)\n', (3793, 3802), False, 'from neural_compressor.experimental import Quantization, common\n'), ((3840, 3866), 'neural_compressor.experimental.common.DataLoader', 'common.DataLoader', (['dataset'], {}), '(dataset)\n', (3857, 3866), False, 'from neural_compressor.experimental import Quantization, common\n'), ((4075, 4106), 'neural_compressor.experimental.Quantization', 'Quantization', (['"""fake_yaml2.yaml"""'], {}), "('fake_yaml2.yaml')\n", (4087, 4106), False, 'from neural_compressor.experimental import Quantization, common\n'), ((4219, 4245), 'neural_compressor.experimental.common.DataLoader', 'common.DataLoader', (['dataset'], {}), '(dataset)\n', (4236, 4245), False, 'from neural_compressor.experimental import Quantization, common\n'), ((4283, 4309), 'neural_compressor.experimental.common.DataLoader', 'common.DataLoader', (['dataset'], {}), '(dataset)\n', (4300, 4309), False, 'from neural_compressor.experimental import Quantization, common\n'), ((1560, 1572), 'tensorflow.Session', 'tf.Session', ([], {}), '()\n', (1570, 1572), True, 'import tensorflow as tf\n'), ((1599, 1655), 'tensorflow.placeholder', 'tf.placeholder', (['tf.float64'], {'shape': '(1, 3, 3, 1)', 'name': '"""x"""'}), "(tf.float64, shape=(1, 3, 3, 1), name='x')\n", (1613, 1655), True, 'import tensorflow as tf\n'), ((1745, 1840), 'tensorflow.nn.conv2d', 'tf.nn.conv2d', ([], {'input': 'x', 'filter': 'y', 'strides': '[1, 1, 1, 1]', 'padding': '"""VALID"""', 'name': '"""op_to_store"""'}), "(input=x, filter=y, strides=[1, 1, 1, 1], padding='VALID', name\n    ='op_to_store')\n", (1757, 1840), True, 'import tensorflow as tf\n'), ((1957, 2045), 'tensorflow.graph_util.convert_variables_to_constants', 'tf.graph_util.convert_variables_to_constants', (['sess', 'sess.graph_def', "['op_to_store']"], {}), "(sess, sess.graph_def, [\n    'op_to_store'])\n", (2001, 2045), True, 'import tensorflow as tf\n'), ((2179, 2218), 'tensorflow.import_graph_def', 'tf.import_graph_def', (['graph_def'], {'name': '""""""'}), "(graph_def, name='')\n", (2198, 2218), True, 'import tensorflow as tf\n'), ((2249, 2259), 'tensorflow.Graph', 'tf.Graph', ([], {}), '()\n', (2257, 2259), True, 'import tensorflow as tf\n'), ((2281, 2304), 'tensorflow.compat.v1.GraphDef', 'tf.compat.v1.GraphDef', ([], {}), '()\n', (2302, 2304), True, 'import tensorflow as tf\n'), ((1685, 1715), 'numpy.random.random', 'np.random.random', (['(2, 2, 1, 1)'], {}), '((2, 2, 1, 1))\n', (1701, 1715), True, 'import numpy as np\n'), ((1892, 1925), 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), '()\n', (1923, 1925), True, 'import tensorflow as tf\n'), ((2319, 2341), 'tensorflow.compat.v1.Session', 'tf.compat.v1.Session', ([], {}), '()\n', (2339, 2341), True, 'import tensorflow as tf\n'), ((2368, 2434), 'tensorflow.compat.v1.placeholder', 'tf.compat.v1.placeholder', (['tf.float64'], {'shape': '(1, 3, 3, 1)', 'name': '"""x"""'}), "(tf.float64, shape=(1, 3, 3, 1), name='x')\n", (2392, 2434), True, 'import tensorflow as tf\n'), ((2534, 2629), 'tensorflow.nn.conv2d', 'tf.nn.conv2d', ([], {'input': 'x', 'filters': 'y', 'strides': '[1, 1, 1, 1]', 'padding': '"""VALID"""', 'name': '"""op_to_store"""'}), "(input=x, filters=y, strides=[1, 1, 1, 1], padding='VALID',\n    name='op_to_store')\n", (2546, 2629), True, 'import tensorflow as tf\n'), ((2757, 2854), 'tensorflow.compat.v1.graph_util.convert_variables_to_constants', 'tf.compat.v1.graph_util.convert_variables_to_constants', (['sess', 'sess.graph_def', "['op_to_store']"], {}), "(sess, sess.graph_def,\n    ['op_to_store'])\n", (2811, 2854), True, 'import tensorflow as tf\n'), ((3057, 3096), 'tensorflow.import_graph_def', 'tf.import_graph_def', (['graph_def'], {'name': '""""""'}), "(graph_def, name='')\n", (3076, 3096), True, 'import tensorflow as tf\n'), ((2474, 2504), 'numpy.random.random', 'np.random.random', (['(2, 2, 1, 1)'], {}), '((2, 2, 1, 1))\n', (2490, 2504), True, 'import numpy as np\n'), ((2682, 2725), 'tensorflow.compat.v1.global_variables_initializer', 'tf.compat.v1.global_variables_initializer', ([], {}), '()\n', (2723, 2725), True, 'import tensorflow as tf\n')]