Example #1
0
def bias_correction_empirical(dataset: tf.data.Dataset):
    """
    Perform bias correction on a given model
    :param dataset: Data passed by user as tf.Dataset type.
    :return: None
    """

    # load a model
    tf.keras.backend.clear_session()
    _ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
    sess = tf.compat.v1.keras.backend.get_session()

    # input parameters for bias correction
    # populate required parameters in two data types QuantParams and BiasCorrectParams

    quant_params = QuantParams(quant_mode='tf_enhanced',
                               round_mode='nearest',
                               use_cuda=True,
                               ops_to_ignore=None)

    bias_correction_params = BiasCorrectionParams(
        batch_size=1,
        num_quant_samples=10,
        num_bias_correct_samples=10,
        input_op_names=['input_1'],
        output_op_names=['fc1000/Softmax'])

    with sess.as_default():
        # run bias correction on the model
        _new_session = BiasCorrection.correct_bias(sess,
                                                   bias_correction_params,
                                                   quant_params, dataset)
    sess.close()
Example #2
0
def bias_correction_after_cle(dataset: tf.data.Dataset):
    """
    Perform bias correction on a given model (mix of empirical and analytical) after
    cross layer equalization.
    :param dataset: Data passed by user as tf.Dataset type.
    :return: None
    """

    # load a model
    tf.keras.backend.clear_session()
    _ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
    sess = tf.compat.v1.keras.backend.get_session()

    # input parameters for bias correction
    # populate required parameters in two data types QuantParams and BiasCorrectParams

    quant_params = QuantParams(quant_mode='tf_enhanced',
                               round_mode='nearest',
                               use_cuda=True,
                               ops_to_ignore=None)

    bias_correction_params = BiasCorrectionParams(
        batch_size=1,
        num_quant_samples=10,
        num_bias_correct_samples=10,
        input_op_names=['input_1'],
        output_op_names=['fc1000/Softmax'])

    with sess.as_default():

        # store conv bns info before performing CLE
        conv_bn_dict = BiasCorrection.find_all_convs_bn_with_activation(
            sess,
            start_op_names=['input_1'],
            output_op_names=['fc1000/Softmax'])

        # perform CLE
        sess_after_cle = equalize_model(sess,
                                        start_op_names=['input_1'],
                                        output_op_names=['fc1000/Softmax'])

        # run empirical and analytical bias correction on the model
        _new_session = BiasCorrection.correct_bias(
            sess_after_cle,
            bias_correction_params,
            quant_params,
            dataset,
            conv_bn_dict=conv_bn_dict,
            perform_only_empirical_bias_corr=False)
    sess.close()
Example #3
0
def bias_correction_single_layer_empirical(dataset: tf.data.Dataset):
    """ perform bias correction on one layer """

    # load a model
    tf.keras.backend.clear_session()
    _ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
    sess = tf.compat.v1.keras.backend.get_session()

    # input parameters for bias correction
    # populate required parameters in two data types QuantParams and BiasCorrectParams

    quant_params = QuantParams(quant_mode='tf_enhanced',
                               round_mode='nearest',
                               use_cuda=True,
                               ops_to_ignore=None)

    bias_correction_params = BiasCorrectionParams(
        batch_size=1,
        num_quant_samples=10,
        num_bias_correct_samples=10,
        input_op_names=['input_1'],
        output_op_names=['fc1000/Softmax'])

    with sess.as_default():
        # initialize model with zero bias
        sess = BiasUtils.initialize_model_with_bias(
            sess, bias_correction_params.input_op_names,
            bias_correction_params.output_op_names)

        # pick a layer for bias correction
        example_conv_layer = sess.graph.get_operation_by_name(
            'res2a_branch2a/Conv2D')

        # invoke bias correction of one layer
        BiasCorrection.bias_correction_per_layer(
            reference_model=sess,
            corrected_model=sess,
            bias_correct_params=bias_correction_params,
            layer_name_to_be_corrected=example_conv_layer.name,
            quant_params=quant_params,
            data_set=dataset)
    sess.close()
    def test_correct_bias_on_mnist_with_analytical_bc(self):
        """  Test correct bias using mnist model and data (analytical bias correction)"""

        def modified_parse(serialized_example):
            """
            Modified mnist dataset parser to provide only image data, without labels
            :param serialized_example:
            :return: Input image and labels
            """
            dim = 28
            features = tf.parse_single_example(serialized_example,
                                               features={'label': tf.FixedLenFeature([], tf.int64),
                                                         'image_raw': tf.FixedLenFeature([], tf.string)})

            # Mnist examples are flattened. Since we aren't performing an augmentations
            # these can remain flattened.
            image = tf.decode_raw(features['image_raw'], tf.uint8)
            image.set_shape([dim*dim])

            # Convert from bytes to floats 0 -> 1.
            image = tf.cast(image, tf.float32) / 255

            return image

        tf.compat.v1.reset_default_graph()
        batch_size = 2
        num_samples = 10
        dataset = tf.data.TFRecordDataset([os.path.join('data', 'mnist', 'validation.tfrecords')]).repeat(1)
        dataset = dataset.map(modified_parse, num_parallel_calls=batch_size)
        dataset = dataset.batch(batch_size=batch_size)

        quant_params = QuantParams()
        bias_correction_params = BiasCorrectionParams(batch_size=batch_size,
                                                      num_quant_samples=num_samples,
                                                      num_bias_correct_samples=num_samples,
                                                      input_op_names=['reshape_input'],
                                                      output_op_names=['dense_1/BiasAdd'])
        sess = load_model_from_meta(meta_path=os.path.join('models', 'mnist_save.meta'))
        BiasCorrection.correct_bias(sess, bias_correction_params, quant_params, dataset,
                                    perform_only_empirical_bias_corr=False)
        self.assertTrue(1)      # Add some actual error check
Example #5
0
def bias_correction_single_layer_analytical():
    """ perform analytical bias correction on one layer """

    # load a model
    tf.keras.backend.clear_session()
    _ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
    sess = tf.compat.v1.keras.backend.get_session()

    # input parameters for bias correction
    # populate required parameters in two data types QuantParams and BiasCorrectParams

    quant_params = QuantParams(quant_mode='tf_enhanced',
                               round_mode='nearest',
                               use_cuda=True,
                               ops_to_ignore=None)

    with sess.as_default():
        # initialize model with zero bias
        sess = BiasUtils.initialize_model_with_bias(sess, ['input_1'],
                                                    ['fc1000/Softmax'])

        # pick a layer for bias correction
        example_conv_layer = sess.graph.get_operation_by_name(
            'res2a_branch2a/Conv2D')

        # get candidate conv bns in the model
        convs_bn_activation_info_dict = BiasCorrection.find_all_convs_bn_with_activation(
            sess, ['input_1'], ['fc1000/Softmax'])

        # make sure to pick example_conv_layer that has a bn op associated with it
        if example_conv_layer in convs_bn_activation_info_dict.keys():

            preceding_bn_layer_info = convs_bn_activation_info_dict[
                example_conv_layer]

            # invoke analytical bias correction on this layer
            BiasCorrection.analytical_bias_correction_per_layer(
                sess, example_conv_layer, preceding_bn_layer_info,
                quant_params)
    sess.close()
Example #6
0
def main(args):
    # configuration for efficient use of gpu
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    print('Loading srgan generator...')
    gen_graph = tf.Graph()
    with gen_graph.as_default():
        gen_sess = tf.Session(config=config, graph=gen_graph)
        with gen_sess.as_default():
            srgan_generator = generator()
            srgan_generator.load_weights(args.weights_path)

    # sort files by filenames, assuming names match in both paths
    lr_images_files = sorted(
        glob.glob(os.path.join(args.images_path, '*LR.png')))
    hr_images_files = sorted(
        glob.glob(os.path.join(args.images_path, '*HR.png')))

    # check if number of images align
    if len(lr_images_files) != len(hr_images_files):
        raise RuntimeError('length of image files doesn`t match,'
                           'need same number of images for both'
                           'low resolution and high resolution!')

    image_files = (lr_images_files, hr_images_files)

    # two list of metrics on all images
    psnr_vals, ssim_vals = evaluate_session(gen_sess, image_files,
                                            srgan_generator.input.name,
                                            srgan_generator.output.name)
    psnr_val = np.mean(psnr_vals)
    ssim_val = np.mean(ssim_vals)
    print(
        f'Mean PSNR and SSIM for given images on original model are: [{psnr_val}, {ssim_val}]'
    )

    # TODO: use a better default dataset for compute encodings when not given by users
    # use low resolution images if no representative lr data are provided

    # use low and high resolution images if no representative lr and hr data are provided
    if args.representative_datapath:
        bc_lr_data = glob.glob(
            os.path.join(args.representative_datapath, '*LR.png'))
        comp_encodings_lr_data = glob.glob(
            os.path.join(args.representative_datapath, '*LR.png'))
        comp_encodings_hr_data = glob.glob(
            os.path.join(args.representative_datapath, '*HR.png'))
    else:
        warnings.warn(
            'No representative input data are given,'
            'bias correction and computation of encodings will be done'
            'on part of all of the low resolution images!')
        bc_lr_data = lr_images_files

        warnings.warn('No representative reference data are given,'
                      'computation of encodings will be done'
                      'on part of all of the high resolution images!')
        comp_encodings_lr_data = lr_images_files
        comp_encodings_hr_data = hr_images_files

    comp_encodings_data = (comp_encodings_lr_data, comp_encodings_hr_data)

    if args.cross_layer_equalization:
        print('Applying cross layer equalization (CLE) to session...')
        gen_sess = equalize_model(
            gen_sess,
            start_op_names=srgan_generator.input.op.name,
            output_op_names=srgan_generator.output.op.name)

    if args.bias_correction:
        print('Applying Bias Correction (BC) to session...')
        # the dataset being evaluated might have varying image sizes
        # so right now only use batch size 1
        batch_size = 1
        num_imgs = len(bc_lr_data)

        quant_params = QuantParams(use_cuda=args.use_cuda,
                                   quant_mode=args.quant_scheme)
        bias_correction_params = BiasCorrectionParams(
            batch_size=batch_size,
            num_quant_samples=min(num_imgs, args.num_quant_samples),
            num_bias_correct_samples=min(num_imgs,
                                         args.num_bias_correct_samples),
            input_op_names=[srgan_generator.input.op.name],
            output_op_names=[srgan_generator.output.op.name])

        ds = make_dataset(bc_lr_data)
        ds = ds.batch(batch_size)

        gen_sess = BiasCorrection.correct_bias(gen_sess,
                                               bias_correction_params,
                                               quant_params, ds)

    # creating quantsim object which inserts quantizer ops
    sim = quantsim.QuantizationSimModel(
        gen_sess,
        starting_op_names=[srgan_generator.input.op.name],
        output_op_names=[srgan_generator.output.op.name],
        quant_scheme=args.quant_scheme,
        default_output_bw=args.default_output_bw,
        default_param_bw=args.default_param_bw)

    # compute activation encodings
    # usually achieves good results when data being used for computing
    # encodings are representative of its task
    partial_eval = partial(evaluate_session,
                           input_name=srgan_generator.input.name,
                           output_name='lambda_3/mul_quantized:0')
    sim.compute_encodings(partial_eval, comp_encodings_data)

    psnr_vals, ssim_vals = evaluate_session(sim.session,
                                            image_files,
                                            srgan_generator.input.name,
                                            'lambda_3/mul_quantized:0',
                                            output_dir=args.output_dir)
    psnr_val = np.mean(psnr_vals)
    ssim_val = np.mean(ssim_vals)

    print(
        f'Mean PSNR and SSIM for given images on quantized model are: [{psnr_val}, {ssim_val}]'
    )
Example #7
0
def run_evaluation(args):
    # Build graph definition
    with tf.Graph().as_default():
        # Create iterator
        tf_records = glob(args.dataset_dir + '/validation*')
        preprocessing_fn = preprocessing_factory.get_preprocessing(
            args.model_name, is_training=False)
        parse_function = wrap_preprocessing(preprocessing_fn,
                                            height=args.image_size,
                                            width=args.image_size,
                                            num_classes=(1001 -
                                                         args.labels_offset),
                                            labels_offset=args.labels_offset)

        dataset = tf.data.TFRecordDataset(tf_records).repeat(1)
        dataset = dataset.map(parse_function, num_parallel_calls=1).apply(
            tf.contrib.data.batch_and_drop_remainder(args.batch_size))
        iterator = dataset.make_initializable_iterator()
        images, labels = iterator.get_next()

        network_fn = nets_factory.get_network_fn(
            args.model_name,
            num_classes=(1001 - args.labels_offset),
            is_training=False)
        with tf.device('/cpu:0'):
            images = tf.placeholder_with_default(images,
                                                 shape=(None, args.image_size,
                                                        args.image_size, 3),
                                                 name='input')
            labels = tf.placeholder_with_default(labels,
                                                 shape=(None, 1001 -
                                                        args.labels_offset),
                                                 name='labels')
        logits, end_points = network_fn(images)
        confidences = tf.nn.softmax(logits, axis=1, name='confidences')
        categorical_preds = tf.argmax(confidences,
                                      axis=1,
                                      name='categorical_preds')
        categorical_labels = tf.argmax(labels,
                                       axis=1,
                                       name='categorical_labels')
        correct_predictions = tf.equal(categorical_labels, categorical_preds)
        top1_acc = tf.reduce_mean(tf.cast(correct_predictions, tf.float32),
                                  name='top1-acc')
        top5_acc = tf.reduce_mean(tf.cast(
            tf.nn.in_top_k(predictions=confidences,
                           targets=tf.cast(categorical_labels, tf.int32),
                           k=5), tf.float32),
                                  name='top5-acc')

        saver = tf.train.Saver()
        sess = tf.Session()

        # Load model from checkpoint
        if not args.ckpt_bn_folded:
            saver.restore(sess, args.checkpoint_path)
        else:
            sess.run(tf.global_variables_initializer())

    # Fold all BatchNorms before QuantSim
    sess, folded_pairs = fold_all_batch_norms(sess, ['IteratorGetNext'],
                                              [logits.name[:-2]])

    if args.ckpt_bn_folded:
        with sess.graph.as_default():
            saver = tf.train.Saver()
            saver.restore(sess, args.checkpoint_path)
    else:
        # Do Cross Layer Equalization and Bias Correction if not loading from a batchnorm folded checkpoint
        sess = equalize_model(sess, ['input'], [logits.op.name])
        conv_bn_dict = BiasCorrection.find_all_convs_bn_with_activation(
            sess, ['input'], [logits.op.name])
        quant_params = QuantParams(quant_mode=args.quant_scheme)
        bias_correction_dataset = tf.data.TFRecordDataset(tf_records).repeat(1)
        bias_correction_dataset = bias_correction_dataset.map(
            lambda x: parse_function(x)[0], num_parallel_calls=1).apply(
                tf.contrib.data.batch_and_drop_remainder(args.batch_size))
        bias_correction_params = BiasCorrectionParams(
            batch_size=args.batch_size,
            num_quant_samples=10,
            num_bias_correct_samples=512,
            input_op_names=['input'],
            output_op_names=[logits.op.name])

        sess = BiasCorrection.correct_bias(
            reference_model=sess,
            bias_correct_params=bias_correction_params,
            quant_params=quant_params,
            data_set=bias_correction_dataset,
            conv_bn_dict=conv_bn_dict,
            perform_only_empirical_bias_corr=True)

    # Define eval_func to use for compute encodings in QuantSim
    def eval_func(session, iterations):
        cnt = 0
        avg_acc_top1 = 0
        session.run('MakeIterator')
        while cnt < iterations or iterations == -1:
            try:
                avg_acc_top1 += session.run('top1-acc:0')
                cnt += 1
            except:
                return avg_acc_top1 / cnt

        return avg_acc_top1 / cnt

    # Select the right quant_scheme
    if args.quant_scheme == 'range_learning_tf':
        quant_scheme = aimet_common.defs.QuantScheme.training_range_learning_with_tf_init
    elif args.quant_scheme == 'range_learning_tf_enhanced':
        quant_scheme = aimet_common.defs.QuantScheme.training_range_learning_with_tf_enhanced_init
    elif args.quant_scheme == 'tf':
        quant_scheme = aimet_common.defs.QuantScheme.post_training_tf
    elif args.quant_scheme == 'tf_enhanced':
        quant_scheme = aimet_common.defs.QuantScheme.post_training_tf_enhanced
    else:
        raise ValueError("Got unrecognized quant_scheme: " + args.quant_scheme)

    # Create QuantizationSimModel
    sim = QuantizationSimModel(
        session=sess,
        starting_op_names=['IteratorGetNext'],
        output_op_names=[logits.name[:-2]],
        quant_scheme=quant_scheme,
        rounding_mode=args.round_mode,
        default_output_bw=args.default_output_bw,
        default_param_bw=args.default_param_bw,
        config_file=args.quantsim_config_file,
    )

    # Run compute_encodings
    sim.compute_encodings(eval_func,
                          forward_pass_callback_args=args.encodings_iterations)

    # Run final evaluation
    sess = sim.session

    top1_acc = eval_func(sess, -1)
    print('Avg accuracy  Top 1: {}'.format(top1_acc))
    def test_analytical_empirical_bias_correction(self):
        """
        Test bn based bias correction hybrid with a user passed in dictionary of conv and bn after cle.
        """

        # create a custom model
        tf.compat.v1.reset_default_graph()
        inputs = tf.keras.Input(shape=(
            32,
            32,
            3,
        ), name="inputs")

        conv_op = tf.keras.layers.Conv2D(
            32, (3, 3),
            kernel_initializer=tf.random_uniform_initializer(-1, 1),
            bias_initializer='random_uniform')(inputs)

        conv1_op = tf.keras.layers.Conv2D(
            32, (3, 3),
            kernel_initializer=tf.random_uniform_initializer(-1, 1),
            bias_initializer='random_uniform')(conv_op)

        bn_op = tf.keras.layers.BatchNormalization(
            fused=True,
            beta_initializer='random_uniform',
            gamma_initializer='random_uniform',
            moving_mean_initializer='random_uniform',
            moving_variance_initializer='random_uniform')(conv1_op,
                                                          training=False)

        conv2_op = tf.keras.layers.Conv2D(
            32, (3, 3),
            kernel_initializer=tf.random_uniform_initializer(-1, 1),
            bias_initializer='random_uniform')(bn_op)

        bn_op2 = tf.keras.layers.BatchNormalization(
            fused=True,
            beta_initializer='random_uniform',
            gamma_initializer='random_uniform',
            moving_mean_initializer='random_uniform',
            moving_variance_initializer='random_uniform')(conv2_op,
                                                          training=False)
        relu_1 = tf.nn.relu(bn_op2)
        conv6_op = tf.keras.layers.Conv2D(32, (3, 3))(relu_1)

        _ = tf.nn.relu(conv6_op)

        init = tf.compat.v1.global_variables_initializer()
        sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
        sess.run(init)

        output_op = sess.graph.get_operation_by_name('Relu_1')
        input_op_name = "inputs"

        bias_corr_input = BiasCorrectionParams(
            batch_size=1,
            num_quant_samples=10,
            num_bias_correct_samples=10,
            input_op_names=[input_op_name],
            output_op_names=[output_op.name])
        quant_params = QuantParams(use_cuda=False)

        np.random.seed(0)

        input_tensor = sess.graph.get_tensor_by_name('inputs:0')
        shape = input_tensor.shape
        dataset = np.random.rand(1, shape[1], shape[2], shape[3])

        # store conv bns info
        conv_bn_dict = BiasCorrection.find_all_convs_bn_with_activation(
            sess, [input_op_name], [output_op.name])

        # perform CLE
        new_sess = equalize_model(sess, input_op_name, output_op.name)
        conv_with_bn_op = new_sess.graph.get_operation_by_name(
            'conv2d_1/Conv2D')
        old_bias_as_numpy = BiasUtils.get_bias_as_numpy_data(
            new_sess, conv_with_bn_op)

        # perform bias correction and check analytical is performed.
        with unittest.mock.patch(
                'aimet_tensorflow.bias_correction.iter_first_x'
        ) as iter_first_x:
            iter_first_x.return_value = [dataset]
            with unittest.mock.patch(
                    'aimet_tensorflow.bias_correction.BiasCorrection.analytical_bias_correction_per_layer',
                    return_value=sess
            ) as mocked_analytical_bias_correction_per_layer:
                updated_sess = BiasCorrection.correct_bias(
                    new_sess,
                    bias_corr_input,
                    quant_params,
                    dataset,
                    conv_bn_dict=conv_bn_dict,
                    perform_only_empirical_bias_corr=False)

        self.assertEqual(
            mocked_analytical_bias_correction_per_layer.call_count, 3)

        sess.close()
        new_sess.close()
    def test_bias_correction_single_layer(self):
        """
        Test bias correction for a single layer api
        """

        tf.compat.v1.reset_default_graph()
        config = tf.compat.v1.ConfigProto()
        config.gpu_options.allow_growth = True

        # create a custom model
        inputs = tf.keras.Input(shape=(
            32,
            16,
            3,
        ))
        conv_op = tf.keras.layers.Conv2D(16, (3, 3))(inputs)
        relu_1 = tf.nn.relu(conv_op)
        conv2_op = tf.keras.layers.Conv2D(32, (3, 3))(relu_1)
        relu_2 = tf.nn.relu(conv2_op)

        # global initializer
        init = tf.compat.v1.global_variables_initializer()
        sess = tf.compat.v1.Session(config=config,
                                    graph=tf.compat.v1.get_default_graph())
        sess.run(init)

        # populate conv with dummy bias and weights
        np.random.seed(0)
        conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
        w_shape = WeightTensorUtils.get_tensor_shape(conv_op)
        w_numpy_data = np.random.rand(w_shape[0], w_shape[1], w_shape[2],
                                      w_shape[3])
        b_shape = BiasUtils.get_shape(conv_op)
        b_numpy_data = np.random.rand(b_shape[0])

        WeightTensorUtils.update_tensor_for_op(sess, conv_op, w_numpy_data)
        BiasUtils.update_bias_for_op(sess, conv_op, b_numpy_data)

        # save and load the updated graph after high bias fold update
        n_sess = aimet_tensorflow.utils.graph_saver.save_and_load_graph(
            './test_update', sess)

        output_op = n_sess.graph.get_operation_by_name('Relu_1')
        conv_op = n_sess.graph.get_operation_by_name('conv2d/Conv2D')
        bias_data = BiasUtils.get_bias_as_numpy_data(n_sess, conv_op)

        input_op_name = conv_op.inputs[0].op.name

        bias_corr_input = BiasCorrectionParams(
            batch_size=1,
            num_quant_samples=10,
            num_bias_correct_samples=10,
            input_op_names=[input_op_name],
            output_op_names=[output_op.name])

        quant_params = QuantParams(use_cuda=False)

        np.random.seed(0)
        shape = conv_op.inputs[0].shape
        dataset = np.random.rand(1, shape[1], shape[2], shape[3])

        with unittest.mock.patch(
                'aimet_tensorflow.bias_correction.iter_first_x'
        ) as iter_first_x:
            iter_first_x.return_value = [dataset]
            BiasCorrection.bias_correction_per_layer(
                reference_model=n_sess,
                corrected_model=n_sess,
                bias_correct_params=bias_corr_input,
                layer_name_to_be_corrected=conv_op.name,
                quant_params=quant_params,
                data_set=dataset)

            conv_op = n_sess.graph.get_operation_by_name('conv2d/Conv2D')
            bias_data_updated = BiasUtils.get_bias_as_numpy_data(
                n_sess, conv_op)

        # needs a validation
        self.assertFalse(np.allclose(bias_data, bias_data_updated, atol=1e-4))
        print('Test completed')

        sess.close()
        n_sess.close()
    def test_bias_correction_model_tf_with_no_bias(self):
        """
        Test bias correction for custom model
        """
        tf.compat.v1.reset_default_graph()
        inputs = tf.keras.Input(shape=(
            32,
            32,
            3,
        ))

        conv_op = tf.keras.layers.Conv2D(32, (3, 3), use_bias=False)(inputs)
        relu_1 = tf.nn.relu(conv_op)

        conv2_op = tf.keras.layers.Conv2D(32, (3, 3), use_bias=False)(relu_1)
        relu_2 = tf.nn.relu(conv2_op)

        conv3_op = tf.keras.layers.Conv2D(32, (3, 3), use_bias=False)(relu_2)
        _ = tf.nn.relu(conv3_op)

        init = tf.compat.v1.global_variables_initializer()
        sess = tf.compat.v1.Session()
        sess.run(init)

        # updating random bias and weight for one conv
        np.random.seed(0)
        conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
        w_shape = WeightTensorUtils.get_tensor_shape(conv_op)
        w_shape = WeightTensorUtils.get_tensor_shape(conv_op)
        w_numpy_data = np.random.rand(w_shape[0], w_shape[1], w_shape[2],
                                      w_shape[3])

        # save and load the updated graph after high bias fold update
        n_sess = save_and_load_graph('./test_update', sess)
        conv_op = n_sess.graph.get_operation_by_name('conv2d/Conv2D')

        input_op_name = conv_op.inputs[0].op.name
        output_op = n_sess.graph.get_operation_by_name('Relu_2')

        input_op_names = [input_op_name]
        output_op_names = [output_op.name]

        batch_size = 1
        num_samples = 10

        np.random.seed(0)
        shape = conv_op.inputs[0].shape

        dataset = np.random.rand(10, 1, shape[1], shape[2], shape[3])
        dataset = tf.convert_to_tensor(dataset)
        dataset = tf.data.Dataset.from_tensor_slices(dataset)

        quant_params = QuantParams(quant_mode='tf', use_cuda=False)
        bias_correction_params = BiasCorrectionParams(
            batch_size=batch_size,
            num_quant_samples=num_samples,
            num_bias_correct_samples=num_samples,
            input_op_names=input_op_names,
            output_op_names=output_op_names)

        conv_op = sess.graph.get_operation_by_name('conv2d_1/Conv2D')
        assert (BiasUtils.is_bias_none(conv_op))
        new_sess = BiasCorrection.correct_bias(
            n_sess,
            bias_correction_params,
            quant_params,
            dataset,
            perform_only_empirical_bias_corr=False)

        conv_op = new_sess.graph.get_operation_by_name('conv2d_1/Conv2D')
        assert (not BiasUtils.is_bias_none(conv_op))

        sess.close()
        n_sess.close()
        new_sess.close()
    def test_analytical_empirical_bias_correction(self):
        """
        Test bn based bias correction for a single layer api
        """

        # create a custom model
        tf.compat.v1.reset_default_graph()
        inputs = tf.keras.Input(shape=(
            32,
            32,
            3,
        ), name="inputs")

        conv1_op = tf.keras.layers.Conv2D(
            32, (3, 3),
            kernel_initializer=tf.random_uniform_initializer(-1, 1),
            bias_initializer='random_uniform')(inputs)

        bn_op = tf.keras.layers.BatchNormalization(
            fused=True,
            beta_initializer='random_uniform',
            gamma_initializer='random_uniform',
            moving_mean_initializer='random_uniform',
            moving_variance_initializer='random_uniform')(conv1_op,
                                                          training=False)

        conv2_op = tf.keras.layers.Conv2D(32, (3, 3))(bn_op)
        conv3_op = tf.keras.layers.Conv2D(32, (3, 3))(conv2_op)
        conv4_op = tf.keras.layers.Conv2D(32, (3, 3))(conv3_op)
        conv5_op = tf.keras.layers.Conv2D(32, (3, 3))(conv4_op)
        bn_op2 = tf.keras.layers.BatchNormalization(
            fused=True,
            beta_initializer='random_uniform',
            gamma_initializer='random_uniform',
            moving_mean_initializer='random_uniform',
            moving_variance_initializer='random_uniform')(conv5_op,
                                                          training=False)
        relu_1 = tf.nn.relu(bn_op2)

        conv6_op = tf.keras.layers.Conv2D(32, (3, 3))(relu_1)

        _ = tf.nn.relu(conv6_op)

        init = tf.compat.v1.global_variables_initializer()
        sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
        sess.run(init)

        output_op = sess.graph.get_operation_by_name('Relu_1')
        input_op_name = "inputs"

        bias_corr_input = BiasCorrectionParams(
            batch_size=1,
            num_quant_samples=10,
            num_bias_correct_samples=10,
            input_op_names=[input_op_name],
            output_op_names=[output_op.name])
        quant_params = QuantParams(use_cuda=False)

        np.random.seed(0)

        input_tensor = sess.graph.get_tensor_by_name('inputs:0')
        shape = input_tensor.shape
        dataset = np.random.rand(1, shape[1], shape[2], shape[3])

        with unittest.mock.patch(
                'aimet_tensorflow.bias_correction.iter_first_x'
        ) as iter_first_x:
            iter_first_x.return_value = [dataset]
            with unittest.mock.patch(
                    'aimet_tensorflow.bias_correction.BiasCorrection.analytical_bias_correction_per_layer',
                    return_value=sess
            ) as mocked_analytical_bias_correction_per_layer:
                with unittest.mock.patch(
                        'aimet_tensorflow.bias_correction.BiasCorrection.bias_correction_per_layer',
                        return_value=sess) as mocked_bias_correction_per_layer:
                    updated_sess = BiasCorrection.correct_bias(
                        sess,
                        bias_corr_input,
                        quant_params,
                        dataset,
                        perform_only_empirical_bias_corr=False)

        assert mocked_bias_correction_per_layer.called
        assert mocked_analytical_bias_correction_per_layer.called
        self.assertEqual(
            mocked_analytical_bias_correction_per_layer.call_count, 3)
        self.assertEqual(mocked_bias_correction_per_layer.call_count,
                         3)  # conv 3,4,5

        sess.close()
        updated_sess.close()
    def test_bn_based_bias_correction_single_layer_functions_invoked(self):
        """
        Test bn based bias correction for a single layer api methods invoked correctly
        """

        # create a custom model
        tf.compat.v1.reset_default_graph()
        inputs = tf.keras.Input(shape=(
            32,
            32,
            3,
        ), name="inputs")
        conv1_op = tf.keras.layers.Conv2D(
            32, (3, 3),
            kernel_initializer=tf.random_uniform_initializer(-1, 1),
            bias_initializer='random_uniform')(inputs)
        conv2_op = tf.keras.layers.Conv2D(32, (3, 3))(conv1_op)
        _ = tf.nn.relu(conv2_op)

        init = tf.compat.v1.global_variables_initializer()
        sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
        sess.run(init)

        conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
        output_op = sess.graph.get_operation_by_name('Relu')
        input_op_name = "inputs"

        bias_corr_input = BiasCorrectionParams(
            batch_size=1,
            num_quant_samples=10,
            num_bias_correct_samples=10,
            input_op_names=[input_op_name],
            output_op_names=[output_op.name])
        quant_params = QuantParams(use_cuda=False)

        np.random.seed(0)
        shape = conv_op.inputs[0].shape
        dataset = np.random.rand(1, shape[1], shape[2], shape[3])

        bias_tensor, weight = BiasCorrection._get_conv_linear_params(
            sess, conv_op)
        q_weight = BiasCorrection._get_quantized_weights(weight, quant_params)

        with unittest.mock.patch(
                'aimet_tensorflow.bias_correction.iter_first_x'
        ) as iter_first_x:
            iter_first_x.return_value = [dataset]
            with unittest.mock.patch(
                    'aimet_tensorflow.bias_correction.BiasCorrection.analytical_bias_correction_per_layer',
                    return_value=sess
            ) as mocked_analytical_bias_correction_per_layer:
                with unittest.mock.patch(
                        'aimet_tensorflow.bias_correction.BiasCorrection.bias_correction_per_layer',
                        return_value=sess) as mocked_bias_correction_per_layer:
                    updated_sess = BiasCorrection.correct_bias(
                        sess,
                        bias_corr_input,
                        quant_params,
                        dataset,
                        perform_only_empirical_bias_corr=False)

        # check if api(s) are invoked
        assert mocked_analytical_bias_correction_per_layer.called
        called_args = mocked_analytical_bias_correction_per_layer.call_args
        assert (called_args[1]['is_first_conv'] == True)
        assert mocked_bias_correction_per_layer.called

        sess.close()
        updated_sess.close()
    def test_depthwise_custom(self):
        """ test depthwise conv2d layer withput bias """

        tf.compat.v1.reset_default_graph()
        inputs = tf.keras.Input(shape=(
            10,
            10,
            3,
        ))
        x = tf.keras.layers.Conv2D(10, (1, 1))(inputs)
        with tf.compat.v1.variable_scope("standalone_depthwise"):
            x = tf.compat.v1.nn.depthwise_conv2d_native(
                x,
                tf.compat.v1.get_variable(
                    initializer=tf.random.truncated_normal(shape=(3, 3, 10,
                                                                  1)),
                    name="depthwise_kernel"), [1, 1, 1, 1], 'VALID')
        _ = tf.nn.relu(x)

        init = tf.compat.v1.global_variables_initializer()
        sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
        sess.run(init)

        op_list = sess.graph.get_operations()
        depthwise_conv_op = sess.graph.get_operation_by_name(
            'standalone_depthwise/DepthwiseConv2dNative')
        input_op = sess.graph.get_operation_by_name('input_1')
        output_op = sess.graph.get_operation_by_name('Relu')

        input_op_names = ['input_1']
        output_op_names = [output_op.name]

        batch_size = 1
        num_samples = 10

        np.random.seed(0)
        shape = input_op.outputs[0].shape

        dataset = np.random.rand(10, 1, shape[1], shape[2], shape[3])
        dataset = tf.convert_to_tensor(dataset)
        dataset = tf.data.Dataset.from_tensor_slices(dataset)

        quant_params = QuantParams(use_cuda=False)
        bias_correction_params = BiasCorrectionParams(
            batch_size=batch_size,
            num_quant_samples=num_samples,
            num_bias_correct_samples=num_samples,
            input_op_names=input_op_names,
            output_op_names=output_op_names)

        assert (BiasUtils.is_bias_none(depthwise_conv_op))

        new_sess = BiasCorrection.correct_bias(sess, bias_correction_params,
                                               quant_params, dataset)

        updated_conv_op = new_sess.graph.get_operation_by_name(
            'standalone_depthwise/DepthwiseConv2dNative')

        assert (not BiasUtils.is_bias_none(updated_conv_op))

        sess.close()
        new_sess.close()
    def test_bias_update_to_dense(self):
        """
        test bias correction on matmul layer
        :return:
        """
        tf.compat.v1.reset_default_graph()

        inputs = tf.keras.Input(shape=(
            32,
            32,
            3,
        ))
        x = tf.keras.layers.Flatten()(inputs)
        dense = tf.keras.layers.Dense(2,
                                      use_bias=False,
                                      activation=tf.nn.softmax,
                                      name="single_residual")(x)
        _ = tf.nn.relu(dense)

        init = tf.compat.v1.global_variables_initializer()
        sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
        sess.run(init)

        op_list = sess.graph.get_operations()

        input_op = sess.graph.get_operation_by_name('input_1')
        output_op = sess.graph.get_operation_by_name('Relu')

        input_op_names = ['input_1']
        output_op_names = [output_op.name]

        op_list = sess.graph.get_operations()

        batch_size = 1
        num_samples = 10

        np.random.seed(0)
        shape = input_op.outputs[0].shape

        dataset = np.random.rand(10, 1, shape[1], shape[2], shape[3])
        dataset = tf.convert_to_tensor(dataset)
        dataset = tf.data.Dataset.from_tensor_slices(dataset)

        quant_params = QuantParams(use_cuda=False)
        bias_correction_params = BiasCorrectionParams(
            batch_size=batch_size,
            num_quant_samples=num_samples,
            num_bias_correct_samples=num_samples,
            input_op_names=input_op_names,
            output_op_names=output_op_names)

        dense_conv_op = sess.graph.get_operation_by_name(
            'single_residual/MatMul')
        assert (BiasUtils.is_bias_none(dense_conv_op))

        new_sess = BiasCorrection.correct_bias(sess, bias_correction_params,
                                               quant_params, dataset)
        updated_dense_conv_op = new_sess.graph.get_operation_by_name(
            'single_residual/MatMul')
        bias = BiasUtils.get_bias_as_numpy_data(new_sess,
                                                updated_dense_conv_op)
        assert (not BiasUtils.is_bias_none(updated_dense_conv_op))

        sess.close()
        new_sess.close()