jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_horovod_syncbn_gpu(self):
        """Test that the SyncBatchNormalization implementation is correct on GPU."""
        # Only do this test if there are GPUs available.
        if not tf.test.is_gpu_available(cuda_only=True):
            self.skipTest(("No GPUs available"))

        hvd.init()
        with tf.device("/gpu:%d" % hvd.local_rank()):
            x_list = [
                tf.convert_to_tensor(np.stack([
                    np.array([
                        [r, r + 1],
                        [r * 2, r * 2 + 1],
                        [r * 3, r * 3 + 1],
                        [r * 4, r * 4 + 1]
                    ], dtype=np.float32)
                    for r in range(hvd.size())
                ]), np.float32),
                tf.convert_to_tensor(np.stack([
                    np.array([
                        [r + 1],
                        [r * 2 + 1],
                        [r * 3 + 1],
                        [r * 4 + 1]
                    ], dtype=np.float32)
                    for r in range(hvd.size())
                ]), np.float32),
            ]

            for x in x_list:
                bn = tf.keras.layers.BatchNormalization(axis=1, fused=False)
                sync_bn = hvd.SyncBatchNormalization(axis=1)
                bn_func = bn(x, training=True)
                sync_bn_func = sync_bn(tf.expand_dims(x[hvd.rank()], 0), training=True)

                try:
                  init = tf.global_variables_initializer()
                except AttributeError:
                  init = tf.compat.v1.global_variables_initializer()
                self.evaluate(init)
                bn_out = self.evaluate(bn_func)
                sync_bn_out = self.evaluate(sync_bn_func)

                self.assertAllClose(sync_bn_out, np.expand_dims(bn_out[hvd.rank()], 0))
                self.assertAllClose(self.evaluate(sync_bn.moving_mean), self.evaluate(bn.moving_mean))
                self.assertAllClose(self.evaluate(sync_bn.moving_variance), self.evaluate(bn.moving_variance))
//  [python]
def test_create_model(self, use_t2t_decoder):
    """Test creating LaserTagger_AR model."""
    input_ids = tf.constant([[0, 2, 3, 1, 0]], dtype=tf.int64)
    input_mask = tf.constant([[1, 1, 1, 1, 0]], dtype=tf.int64)
    segment_ids = tf.constant([[0, 0, 0, 0, 0]], dtype=tf.int64)
    labels = tf.constant([[0, 0, 1, 0, 0]], dtype=tf.int64)
    labels_mask = tf.constant([[0, 1, 1, 0, 0]], dtype=tf.int64)

    model_fn_builder = _get_model_builder(use_t2t_decoder)
    (loss, _, pred) = model_fn_builder._create_model(
        tf.estimator.ModeKeys.TRAIN, input_ids, input_mask, segment_ids, labels,
        labels_mask)

    with self.test_session() as sess:
      sess.run(tf.initialize_all_variables())
      out = sess.run({"loss": loss, "pred": pred})

      self.assertEqual(out["loss"].shape, ())
      self.assertEqual(out["pred"].shape, labels.shape)
//  [python]
def test_create_model(self, use_t2t_decoder):
    """Test creating LaserTagger_AR model."""
    input_ids = tf.constant([[0, 2, 3, 1, 0]], dtype=tf.int64)
    input_mask = tf.constant([[1, 1, 1, 1, 0]], dtype=tf.int64)
    segment_ids = tf.constant([[0, 0, 0, 0, 0]], dtype=tf.int64)
    labels = tf.constant([[0, 0, 1, 0, 0]], dtype=tf.int64)
    labels_mask = tf.constant([[0, 1, 1, 0, 0]], dtype=tf.int64)

    model_fn_builder = _get_model_builder(use_t2t_decoder)
    (loss, _, pred) = model_fn_builder._create_model(
        tf.estimator.ModeKeys.TRAIN, input_ids, input_mask, segment_ids, labels,
        labels_mask)

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      out = sess.run({"loss": loss, "pred": pred})

      self.assertEqual(out["loss"].shape, ())
      self.assertEqual(out["pred"].shape, labels.shape)
//  [python]
def test_horovod_syncbn_cpu(self):
        """Test that the SyncBatchNormalization implementation is correct on CPU."""

        hvd.init()
        with tf.device("/cpu:0"):
            x_list = [
                tf.convert_to_tensor(np.stack([
                    np.array([
                        [r, r + 1],
                        [r * 2, r * 2 + 1],
                        [r * 3, r * 3 + 1],
                        [r * 4, r * 4 + 1]
                    ], dtype=np.float32)
                    for r in range(hvd.size())
                ]), np.float32),
                tf.convert_to_tensor(np.stack([
                    np.array([
                        [r + 1],
                        [r * 2 + 1],
                        [r * 3 + 1],
                        [r * 4 + 1]
                    ], dtype=np.float32)
                    for r in range(hvd.size())
                ]), np.float32),
            ]

            for x in x_list:
                bn = tf.keras.layers.BatchNormalization(axis=1, fused=False)
                sync_bn = hvd.SyncBatchNormalization(axis=1)
                bn_func = bn(x, training=True)
                sync_bn_func = sync_bn(tf.expand_dims(x[hvd.rank()], 0), training=True)

                try:
                  init = tf.compat.v1.global_variables_initializer()
                except AttributeError:
                  init = tf.compat.v1.global_variables_initializer()
                self.evaluate(init)
                bn_out = self.evaluate(bn_func)
                sync_bn_out = self.evaluate(sync_bn_func)

                self.assertAllClose(sync_bn_out, np.expand_dims(bn_out[hvd.rank()], 0))
                self.assertAllClose(self.evaluate(sync_bn.moving_mean), self.evaluate(bn.moving_mean))
                self.assertAllClose(self.evaluate(sync_bn.moving_variance), self.evaluate(bn.moving_variance))
//  [python]
def test_horovod_syncbn_cpu(self):
        """Test that the SyncBatchNormalization implementation is correct on CPU."""

        hvd.init()
        with tf.device("/cpu:0"):
            x_list = [
                tf.convert_to_tensor(np.stack([
                    np.array([
                        [r, r + 1],
                        [r * 2, r * 2 + 1],
                        [r * 3, r * 3 + 1],
                        [r * 4, r * 4 + 1]
                    ], dtype=np.float32)
                    for r in range(hvd.size())
                ]), np.float32),
                tf.convert_to_tensor(np.stack([
                    np.array([
                        [r + 1],
                        [r * 2 + 1],
                        [r * 3 + 1],
                        [r * 4 + 1]
                    ], dtype=np.float32)
                    for r in range(hvd.size())
                ]), np.float32),
            ]

            for x in x_list:
                bn = tf.keras.layers.BatchNormalization(axis=1, fused=False)
                sync_bn = hvd.SyncBatchNormalization(axis=1)
                bn_func = bn(x, training=True)
                sync_bn_func = sync_bn(tf.expand_dims(x[hvd.rank()], 0), training=True)

                try:
                  init = tf.global_variables_initializer()
                except AttributeError:
                  init = tf.compat.v1.global_variables_initializer()
                self.evaluate(init)
                bn_out = self.evaluate(bn_func)
                sync_bn_out = self.evaluate(sync_bn_func)

                self.assertAllClose(sync_bn_out, np.expand_dims(bn_out[hvd.rank()], 0))
                self.assertAllClose(self.evaluate(sync_bn.moving_mean), self.evaluate(bn.moving_mean))
                self.assertAllClose(self.evaluate(sync_bn.moving_variance), self.evaluate(bn.moving_variance))
//  [python]
def test_model_fn_eval(self):
    with self.session() as sess:
      model_fn_builder = _get_model_builder()
      model_fn = model_fn_builder.build()
      output_spec = model_fn(
          self._features,
          labels=None,
          mode=tf.estimator.ModeKeys.EVAL,
          params=None)
      metric_fn = output_spec.eval_metrics[0]
      metric_fn_args = output_spec.eval_metrics[1]
      self.assertLen(metric_fn_args, 4)
      metrics = metric_fn(*metric_fn_args)
      sess.run([tf.global_variables_initializer(),
                tf.local_variables_initializer()])

      def check_metric_shape(metric):
        val_node, update_op = metric
        sess.run(update_op)
        val = sess.run(val_node)
        self.assertAllEqual(val.shape, [])

      self.assertLen(metrics, 2)
      check_metric_shape(metrics["eval_loss"])
      check_metric_shape(metrics["sentence_level_acc"])
//  [python]
def main(_):
    # Load Data
    raw_data = reader.load_data(FLAGS.parse_data_path)
    train_sents, train_trees, dev_sents, dev_trees, vocab_dict, pos_dict, label_dict, feature_tpl = raw_data # items in ids
    # Load Config
    config_dict = load_config(FLAGS.parse_model_config_path)
    config = get_config(config_dict, FLAGS.parse_lang)
    model_var_scope = get_model_var_scope(FLAGS.parse_scope_name, FLAGS.parse_lang)
    print ("NOTICE: Parsing model variable scope is %s" % model_var_scope)
    with tf.Session() as session:
        with tf.variable_scope(model_var_scope, reuse=True):
            m = NNParser(config=config)
    
    # CheckPoint State
    if not os.path.exists(FLAGS.parse_train_dir):
        os.makedirs(FLAGS.parse_train_dir)

    ckpt = tf.train.get_checkpoint_state(FLAGS.parse_train_dir)
    if ckpt:
        print("Loading model parameters from %s" % ckpt.model_checkpoint_path)
        m.saver.restore(session, tf.train.latest_checkpoint(FLAGS.parse_train_dir))
    else:
        print("Created model with fresh parameters.")
        session.run(tf.initialize_all_variables())
    
    # train dataset should be generated only once and called by run_epoch function
    # max_tree_num = 8000 # for memory saving 
    for i in range(config.max_max_epoch):
        lr_decay = config.lr_decay ** max(i - config.max_epoch, 0.0)
        m.assign_lr(session, config.learning_rate * lr_decay)
        print("Epoch: %d Learning rate: %.4f" % (i + 1, session.run(m.lr)))
        # Training: Check training set perplexity
        instance_path = os.path.join(FLAGS.parse_data_path, "instance")  # saved example instance path
        train_X, train_Y = transition_system.generate_examples(train_sents, train_trees, label_dict, feature_tpl, instance_path, is_train = True)
        train_iter = transition_system.iter_examples(train_X, train_Y, m.batch_size)
        train_perplexity = run_epoch(session, m, m.train_op, train_iter)
        print("Epoch: %d Train Perplexity: %.3f" % (i + 1, train_perplexity))
        # Dev: Check development set perplexity
        dev_X, dev_Y = transition_system.generate_examples(dev_sents, dev_trees, label_dict, feature_tpl, instance_path, is_train = False)
        dev_iter = transition_system.iter_examples(dev_X, dev_Y, m.batch_size)
        dev_perplexity = run_epoch(session, m, tf.no_op(), dev_iter, is_train = False)
        print("Epoch: %d Dev Perplexity: %.3f" % (i + 1, dev_perplexity))
//  [python]
def main(_):
    # Load Data
    raw_data = reader.load_data(FLAGS.parse_data_path)
    train_sents, train_trees, dev_sents, dev_trees, vocab_dict, pos_dict, label_dict, feature_tpl = raw_data # items in ids
    # Load Config
    config_dict = load_config(FLAGS.parse_model_config_path)
    config = get_config(config_dict, FLAGS.parse_lang)
    model_var_scope = get_model_var_scope(FLAGS.parse_scope_name, FLAGS.parse_lang)
    print ("NOTICE: Parsing model variable scope is %s" % model_var_scope)
    with tf.Session() as session:
        with tf.variable_scope(model_var_scope, reuse=True):
            m = NNParser(config=config)
    
    # CheckPoint State
    if not os.path.exists(FLAGS.parse_train_dir):
        os.makedirs(FLAGS.parse_train_dir)

    ckpt = tf.train.get_checkpoint_state(FLAGS.parse_train_dir)
    if ckpt:
        print("Loading model parameters from %s" % ckpt.model_checkpoint_path)
        m.saver.restore(session, tf.train.latest_checkpoint(FLAGS.parse_train_dir))
    else:
        print("Created model with fresh parameters.")
        session.run(tf.global_variables_initializer())
    
    # train dataset should be generated only once and called by run_epoch function
    # max_tree_num = 8000 # for memory saving 
    for i in range(config.max_max_epoch):
        lr_decay = config.lr_decay ** max(i - config.max_epoch, 0.0)
        m.assign_lr(session, config.learning_rate * lr_decay)
        print("Epoch: %d Learning rate: %.4f" % (i + 1, session.run(m.lr)))
        # Training: Check training set perplexity
        instance_path = os.path.join(FLAGS.parse_data_path, "instance")  # saved example instance path
        train_X, train_Y = transition_system.generate_examples(train_sents, train_trees, label_dict, feature_tpl, instance_path, is_train = True)
        train_iter = transition_system.iter_examples(train_X, train_Y, m.batch_size)
        train_perplexity = run_epoch(session, m, m.train_op, train_iter)
        print("Epoch: %d Train Perplexity: %.3f" % (i + 1, train_perplexity))
        # Dev: Check development set perplexity
        dev_X, dev_Y = transition_system.generate_examples(dev_sents, dev_trees, label_dict, feature_tpl, instance_path, is_train = False)
        dev_iter = transition_system.iter_examples(dev_X, dev_Y, m.batch_size)
        dev_perplexity = run_epoch(session, m, tf.no_op(), dev_iter, is_train = False)
        print("Epoch: %d Dev Perplexity: %.3f" % (i + 1, dev_perplexity))
//  [python]
def load_model():
    params = bts_parameters(
        encoder=args.encoder,
        height=args.input_height,
        width=args.input_width,
        batch_size=None,
        dataset=None,
        max_depth=args.max_depth,
        num_gpus=None,
        num_threads=None,
        num_epochs=None,
    )
    
    model = BtsModel(params, 'test', image, None, focal=focals, bn_training=False)
    
    # INIT
    sess.run(tf.initialize_all_variables())
    sess.run(tf.local_variables_initializer())
    coordinator = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coordinator)

    # SAVER
    train_saver = tf.train.Saver()

    with tf.device('/cpu:0'):
        restore_path = args.checkpoint_path
        # RESTORE
        train_saver.restore(sess, restore_path)
    
    return model
//  [python]
def load_model():
    params = bts_parameters(
        encoder=args.encoder,
        height=args.input_height,
        width=args.input_width,
        batch_size=None,
        dataset=None,
        max_depth=args.max_depth,
        num_gpus=None,
        num_threads=None,
        num_epochs=None,
    )
    
    model = BtsModel(params, 'test', image, None, focal=focals, bn_training=False)
    
    # INIT
    sess.run(tf.global_variables_initializer())
    sess.run(tf.local_variables_initializer())
    coordinator = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coordinator)

    # SAVER
    train_saver = tf.train.Saver()

    with tf.device('/cpu:0'):
        restore_path = args.checkpoint_path
        # RESTORE
        train_saver.restore(sess, restore_path)
    
    return model
//  [python]
def convert_to_coverage_model(self):
    """Load non-coverage checkpoint, add initialized extra variables for coverage, and save as new checkpoint"""
    tf.logging.info("converting non-coverage model to coverage model..")

    # initialize an entire coverage model from scratch
    sess = tf.Session(config=util.get_config())
    print("initializing everything...")
    sess.run(tf.initialize_all_variables())

    # load all non-coverage weights from checkpoint
    saver = tf.train.Saver([v for v in tf.global_variables() if "coverage" not in v.name and "Adagrad" not in v.name])
    print("restoring non-coverage variables...")
    curr_ckpt = util.load_ckpt(saver, sess)
    print("restored.")

    # save this model and quit
    new_fname = curr_ckpt + '_cov_init'
    print("saving model to %s..." % (new_fname))
    new_saver = tf.train.Saver() # this one will save all variables that now exist
    new_saver.save(sess, new_fname)
    print("saved.")
    exit()
//  [python]
def convert_to_coverage_model(self):
    """Load non-coverage checkpoint, add initialized extra variables for coverage, and save as new checkpoint"""
    tf.logging.info("converting non-coverage model to coverage model..")

    # initialize an entire coverage model from scratch
    sess = tf.Session(config=util.get_config())
    print("initializing everything...")
    sess.run(tf.global_variables_initializer())

    # load all non-coverage weights from checkpoint
    saver = tf.train.Saver([v for v in tf.global_variables() if "coverage" not in v.name and "Adagrad" not in v.name])
    print("restoring non-coverage variables...")
    curr_ckpt = util.load_ckpt(saver, sess)
    print("restored.")

    # save this model and quit
    new_fname = curr_ckpt + '_cov_init'
    print("saving model to %s..." % (new_fname))
    new_saver = tf.train.Saver() # this one will save all variables that now exist
    new_saver.save(sess, new_fname)
    print("saved.")
    exit()
//  [python]
def convert_to_reinforce_model(self):
    """Load non-reinforce checkpoint, add initialized extra variables for reinforce, and save as new checkpoint"""
    tf.logging.info("converting non-reinforce model to reinforce model..")

    # initialize an entire reinforce model from scratch
    sess = tf.Session(config=util.get_config())
    print("initializing everything...")
    sess.run(tf.initialize_all_variables())

    # load all non-reinforce weights from checkpoint
    saver = tf.train.Saver([v for v in tf.global_variables() if "reinforce" not in v.name and "Adagrad" not in v.name])
    print("restoring non-reinforce variables...")
    curr_ckpt = util.load_ckpt(saver, sess)
    print("restored.")

    # save this model and quit
    new_fname = curr_ckpt + '_rl_init'
    print("saving model to %s..." % (new_fname))
    new_saver = tf.train.Saver() # this one will save all variables that now exist
    new_saver.save(sess, new_fname)
    print("saved.")
    exit()
//  [python]
def convert_to_reinforce_model(self):
    """Load non-reinforce checkpoint, add initialized extra variables for reinforce, and save as new checkpoint"""
    tf.logging.info("converting non-reinforce model to reinforce model..")

    # initialize an entire reinforce model from scratch
    sess = tf.Session(config=util.get_config())
    print("initializing everything...")
    sess.run(tf.global_variables_initializer())

    # load all non-reinforce weights from checkpoint
    saver = tf.train.Saver([v for v in tf.global_variables() if "reinforce" not in v.name and "Adagrad" not in v.name])
    print("restoring non-reinforce variables...")
    curr_ckpt = util.load_ckpt(saver, sess)
    print("restored.")

    # save this model and quit
    new_fname = curr_ckpt + '_rl_init'
    print("saving model to %s..." % (new_fname))
    new_saver = tf.train.Saver() # this one will save all variables that now exist
    new_saver.save(sess, new_fname)
    print("saved.")
    exit()
//  [python]
def test_insert_quant_op_forward_pass_simple_rnn(self):

        """ test insertion of quant ops to recurrent layer with conditional blocks """

        tf.reset_default_graph()
        np.random.seed(0)
        tf.set_random_seed(0)
        sess = tf.Session()

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add an RNN layer
            x = tf.keras.layers.SimpleRNN(12,
                                          kernel_initializer='glorot_uniform',
                                          recurrent_initializer='orthogonal')(inputs)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="simplernn_model")(x)

        init = tf.initialize_all_variables()
        sess.run(init)

        self.validate_simple_rnn_auto_insertion_and_forward_pass(sess)
        sess.close()
//  [python]
def test_insert_quant_op_forward_pass_simple_rnn(self):

        """ test insertion of quant ops to recurrent layer with conditional blocks """

        tf.reset_default_graph()
        np.random.seed(0)
        tf.set_random_seed(0)
        sess = tf.Session()

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add an RNN layer
            x = tf.keras.layers.SimpleRNN(12,
                                          kernel_initializer='glorot_uniform',
                                          recurrent_initializer='orthogonal')(inputs)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="simplernn_model")(x)

        init = tf.global_variables_initializer()
        sess.run(init)

        self.validate_simple_rnn_auto_insertion_and_forward_pass(sess)
        sess.close()
//  [python]
def test_insert_quant_op_forward_pass_simple_rnn_with_relu(self):

        """ test insertion of quant ops to simple rnn with relu """

        tf.reset_default_graph()
        np.random.seed(0)
        tf.set_random_seed(0)
        sess = tf.Session()

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add an RNN layer
            x = tf.keras.layers.SimpleRNN(12, activation='relu',
                                          kernel_initializer='glorot_uniform',
                                          recurrent_initializer='orthogonal')(inputs)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="simplernn_model")(x)

        init = tf.initialize_all_variables()
        sess.run(init)

        self.validate_simple_rnn_auto_insertion_and_forward_pass(sess)
        sess.close()
//  [python]
def test_insert_quant_op_forward_pass_simple_rnn_with_relu(self):

        """ test insertion of quant ops to simple rnn with relu """

        tf.reset_default_graph()
        np.random.seed(0)
        tf.set_random_seed(0)
        sess = tf.Session()

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add an RNN layer
            x = tf.keras.layers.SimpleRNN(12, activation='relu',
                                          kernel_initializer='glorot_uniform',
                                          recurrent_initializer='orthogonal')(inputs)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="simplernn_model")(x)

        init = tf.global_variables_initializer()
        sess.run(init)

        self.validate_simple_rnn_auto_insertion_and_forward_pass(sess)
        sess.close()
//  [python]
def test_insert_quant_op_forward_pass_simple_rnn_multiple_layers(self):

        """ test insertion of quant ops to simple rnn with multiple layes """

        tf.reset_default_graph()
        np.random.seed(0)
        tf.set_random_seed(0)
        sess = tf.Session()

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add an RNN layer
            x = tf.keras.layers.SimpleRNN(12, activation='tanh',
                                          kernel_initializer='glorot_uniform',
                                          recurrent_initializer='orthogonal',
                                          return_sequences=True)(inputs)
            x = tf.keras.layers.SimpleRNN(12, name='rnn1', activation='relu', return_sequences=True)(x)
            x = tf.keras.layers.SimpleRNN(12, name='rnn2', activation='tanh')(x)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="simplernn_model")(x)

        init = tf.initialize_all_variables()
        sess.run(init)

        self.validate_simple_rnn_auto_insertion_and_forward_pass(sess)
        # note - we will need to disable quantizers on identity nodes in this case

        sess.close()
//  [python]
def test_insert_quant_op_forward_pass_simple_rnn_multiple_layers(self):

        """ test insertion of quant ops to simple rnn with multiple layes """

        tf.reset_default_graph()
        np.random.seed(0)
        tf.set_random_seed(0)
        sess = tf.Session()

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add an RNN layer
            x = tf.keras.layers.SimpleRNN(12, activation='tanh',
                                          kernel_initializer='glorot_uniform',
                                          recurrent_initializer='orthogonal',
                                          return_sequences=True)(inputs)
            x = tf.keras.layers.SimpleRNN(12, name='rnn1', activation='relu', return_sequences=True)(x)
            x = tf.keras.layers.SimpleRNN(12, name='rnn2', activation='tanh')(x)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="simplernn_model")(x)

        init = tf.global_variables_initializer()
        sess.run(init)

        self.validate_simple_rnn_auto_insertion_and_forward_pass(sess)
        # note - we will need to disable quantizers on identity nodes in this case

        sess.close()
//  [python]
def test_quantize_lstm_default_quantsim_and_forward_pass(self):
        """ Test connected graph construction on a model with lstm op """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add a LSTM layer with 12 internal units.
            x = tf.keras.layers.LSTM(12)(inputs)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="lstm_model")(x)

        init = tf.initialize_all_variables()
        sess.run(init)
        # _ = tf.summary.FileWriter('./lstm', sess.graph)

        sim = QuantizationSimModel(sess, ['input_1'], ['lstm_model/Softmax'],
                                   use_cuda=False)

        # validate quantsim
        # get ops and make sure we have a quantized op added to the conditional block
        quantized_graph_op_names = self._get_quant_ops_from_tf_graph(sim.session.graph)

        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names)

        # validate forward pass
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim)

        # close tf sessions
        sess.close()
        sim.session.close()
        del sim
//  [python]
def test_quantize_lstm_default_quantsim_and_forward_pass(self):
        """ Test connected graph construction on a model with lstm op """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add a LSTM layer with 12 internal units.
            x = tf.keras.layers.LSTM(12)(inputs)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="lstm_model")(x)

        init = tf.global_variables_initializer()
        sess.run(init)
        # _ = tf.summary.FileWriter('./lstm', sess.graph)

        sim = QuantizationSimModel(sess, ['input_1'], ['lstm_model/Softmax'],
                                   use_cuda=False)

        # validate quantsim
        # get ops and make sure we have a quantized op added to the conditional block
        quantized_graph_op_names = self._get_quant_ops_from_tf_graph(sim.session.graph)

        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names)

        # validate forward pass
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim)

        # close tf sessions
        sess.close()
        sim.session.close()
        del sim
//  [python]
def test_quantize_lstm_sigmoid_quantsim_and_forward_pass(self):
        """ Test connected graph construction on a model with lstm op """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add a LSTM layer with 12 internal units.
            x = tf.keras.layers.LSTM(12, recurrent_activation='sigmoid')(inputs)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="lstm_model")(x)

        init = tf.initialize_all_variables()
        sess.run(init)
        # _ = tf.summary.FileWriter('./lstm', sess.graph)

        sim = QuantizationSimModel(sess, ['input_1'], ['lstm_model/Softmax'],
                                   use_cuda=False)

        # validate quantsim
        # get ops and make sure we have a quantized op added to the conditional block
        quantized_graph_op_names = self._get_quant_ops_from_tf_graph(sim.session.graph)

        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names)

        # validate forward pass
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim)

        # close tf sessions
        sess.close()
        sim.session.close()
        del sim
//  [python]
def test_quantize_lstm_sigmoid_quantsim_and_forward_pass(self):
        """ Test connected graph construction on a model with lstm op """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add a LSTM layer with 12 internal units.
            x = tf.keras.layers.LSTM(12, recurrent_activation='sigmoid')(inputs)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="lstm_model")(x)

        init = tf.global_variables_initializer()
        sess.run(init)
        # _ = tf.summary.FileWriter('./lstm', sess.graph)

        sim = QuantizationSimModel(sess, ['input_1'], ['lstm_model/Softmax'],
                                   use_cuda=False)

        # validate quantsim
        # get ops and make sure we have a quantized op added to the conditional block
        quantized_graph_op_names = self._get_quant_ops_from_tf_graph(sim.session.graph)

        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names)

        # validate forward pass
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim)

        # close tf sessions
        sess.close()
        sim.session.close()
        del sim
//  [python]
def test_quantize_simple_rnn_export(self):
        """ Test model export for recurrent models """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add an RNN layer with 12 internal units.
            x = tf.keras.layers.SimpleRNN(10, name='rnn1', return_sequences=True)(inputs)
            x = tf.keras.layers.SimpleRNN(10, name='rnn2')(x)

            _ = tf.keras.layers.Dense(10, activation=tf.nn.softmax,
                                      name="fc")(x)

            init = tf.global_variables_initializer()
            sess.run(init)

        sim = QuantizationSimModel(sess, ['input_1'], ['fc/Softmax'],
                                   use_cuda=False)

        def dummy_forward_pass(sess, args):
            model_output = sess.graph.get_tensor_by_name('fc/Softmax:0')
            model_input = sess.graph.get_tensor_by_name('input_1:0')
            dummy_input = np.random.randn(1, 3, 100)
            sess.run(model_output, feed_dict={model_input: dummy_input})

        sim.compute_encodings(dummy_forward_pass, None)
        sim.export('./data', 'rnn_quantsim')

        new_sess = load_model_from_meta('./data/rnn_quantsim.meta')

        dummy_forward_pass(new_sess, None)

        all_op_types = [op.type for op in new_sess.graph.get_operations()]
        self.assertNotIn('QcQuantize', all_op_types)
        self.assertNotIn('QcQuantizeRecurrentParam', all_op_types)

        # Load the encodings file to check if the encodings were exported correctly
        with open("./data/rnn_quantsim.encodings", "r") as encodings_file:
            encodings = json.load(encodings_file)
            self.assertEqual(8, len(encodings['activation_encodings']))
            self.assertEqual(5, len(encodings['param_encodings']))

        # close tf sessions
        sess.close()
        sim.session.close()
        del sim
//  [python]
def test_quantize_simple_rnn_save_and_load_checkpoint(self):
        """ Test model export for recurrent models """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add an RNN layer with 12 internal units.
            x = tf.keras.layers.SimpleRNN(10, name='rnn1', return_sequences=True)(inputs)
            x = tf.keras.layers.SimpleRNN(10, name='rnn2')(x)

            _ = tf.keras.layers.Dense(10, activation=tf.nn.softmax,
                                      name="fc")(x)

            init = tf.initialize_all_variables()
            sess.run(init)

        sim = QuantizationSimModel(sess, ['input_1'], ['fc/Softmax'],
                                   use_cuda=False)

        def eval(sess, input_tensor):
            model_output = sess.graph.get_tensor_by_name('fc/Softmax:0')
            model_input = sess.graph.get_tensor_by_name('input_1:0')
            out = sess.run(model_output, feed_dict={model_input: input_tensor})
            return out

        def dummy_forward_pass(sess, args):
            dummy_input = np.random.randn(1, 3, 100)
            eval(sess, dummy_input)

        sim.compute_encodings(dummy_forward_pass, None)
        random_tensor = np.random.randn(1, 3, 100)
        old_out = eval(sim.session, random_tensor)

        save_checkpoint(sim, './data/', 'simple_rnn_save')
        new_sim = load_checkpoint('./data', 'simple_rnn_save')

        # Check to make sure that inference through the new sim produces exactly the same output as the old sim
        # This checks that quantization parameters have been restored correctly
        # Also checks that we are able to invoke quantize-dequantize ops in the new session (so pymo objects were
        # restored correctly etc.)
        new_out = eval(new_sim.session, random_tensor)
        self.assertTrue(np.allclose(old_out, new_out))
        sim.session.close()
        del sim
//  [python]
def test_quantize_simple_rnn_save_and_load_checkpoint(self):
        """ Test model export for recurrent models """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add an RNN layer with 12 internal units.
            x = tf.keras.layers.SimpleRNN(10, name='rnn1', return_sequences=True)(inputs)
            x = tf.keras.layers.SimpleRNN(10, name='rnn2')(x)

            _ = tf.keras.layers.Dense(10, activation=tf.nn.softmax,
                                      name="fc")(x)

            init = tf.global_variables_initializer()
            sess.run(init)

        sim = QuantizationSimModel(sess, ['input_1'], ['fc/Softmax'],
                                   use_cuda=False)

        def eval(sess, input_tensor):
            model_output = sess.graph.get_tensor_by_name('fc/Softmax:0')
            model_input = sess.graph.get_tensor_by_name('input_1:0')
            out = sess.run(model_output, feed_dict={model_input: input_tensor})
            return out

        def dummy_forward_pass(sess, args):
            dummy_input = np.random.randn(1, 3, 100)
            eval(sess, dummy_input)

        sim.compute_encodings(dummy_forward_pass, None)
        random_tensor = np.random.randn(1, 3, 100)
        old_out = eval(sim.session, random_tensor)

        save_checkpoint(sim, './data/', 'simple_rnn_save')
        new_sim = load_checkpoint('./data', 'simple_rnn_save')

        # Check to make sure that inference through the new sim produces exactly the same output as the old sim
        # This checks that quantization parameters have been restored correctly
        # Also checks that we are able to invoke quantize-dequantize ops in the new session (so pymo objects were
        # restored correctly etc.)
        new_out = eval(new_sim.session, random_tensor)
        self.assertTrue(np.allclose(old_out, new_out))
        sim.session.close()
        del sim
//  [python]
def test_quantize_lstm_time_major_true_quantsim_and_forward_pass(self):
        """ Test connected graph construction on a model with lstm op """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add a LSTM layer with 12 internal units.
            x = tf.keras.layers.LSTM(12, time_major=True, name='lstm_tm')(inputs)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="lstm_model")(x)

        init = tf.initialize_all_variables()
        sess.run(init)
        # _ = tf.summary.FileWriter('./lstm', sess.graph)

        sim = QuantizationSimModel(sess, ['input_1'], ['lstm_model/Softmax'],
                                   use_cuda=False)

        # validate quantsim
        # get ops and make sure we have a quantized op added to the conditional blocks
        quantized_graph_op_names = self._get_quant_ops_from_tf_graph(sim.session.graph)

        batches = 32
        def dummy_forward_pass(sess, args):
            model_output = sess.graph.get_tensor_by_name('lstm_model/Softmax:0')
            model_input = sess.graph.get_tensor_by_name('input_1:0')
            dummy_input = np.random.randn(batches, 3, 100)
            sess.run(model_output, feed_dict={model_input: dummy_input})

        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names, 'lstm_tm')

        # validate forward pass
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim)

        # close tf sessions
        sess.close()
        sim.session.close()
        del sim
//  [python]
def test_quantize_lstm_time_major_true_quantsim_and_forward_pass(self):
        """ Test connected graph construction on a model with lstm op """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add a LSTM layer with 12 internal units.
            x = tf.keras.layers.LSTM(12, time_major=True, name='lstm_tm')(inputs)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="lstm_model")(x)

        init = tf.global_variables_initializer()
        sess.run(init)
        # _ = tf.summary.FileWriter('./lstm', sess.graph)

        sim = QuantizationSimModel(sess, ['input_1'], ['lstm_model/Softmax'],
                                   use_cuda=False)

        # validate quantsim
        # get ops and make sure we have a quantized op added to the conditional blocks
        quantized_graph_op_names = self._get_quant_ops_from_tf_graph(sim.session.graph)

        batches = 32
        def dummy_forward_pass(sess, args):
            model_output = sess.graph.get_tensor_by_name('lstm_model/Softmax:0')
            model_input = sess.graph.get_tensor_by_name('input_1:0')
            dummy_input = np.random.randn(batches, 3, 100)
            sess.run(model_output, feed_dict={model_input: dummy_input})

        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names, 'lstm_tm')

        # validate forward pass
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim)

        # close tf sessions
        sess.close()
        sim.session.close()
        del sim
//  [python]
def test_quantize_lstm_deepspeech_time_major_true_quantsim_and_forward_pass(self):
        """ Test connected graph construction on a model with lstm op """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add a LSTM layer with 12 internal units.
            x = tf.keras.layers.LSTM(12,
                                     unit_forget_bias=False,
                                     time_major=True,
                                     return_sequences=True,
                                     name='lstm_stacked')(inputs)
            x2 = tf.keras.layers.LSTM(12, name='last_lstm')(x)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="lstm_model")(x2)

        init = tf.initialize_all_variables()
        sess.run(init)
        # _ = tf.summary.FileWriter('./lstm', sess.graph)

        sim = QuantizationSimModel(sess, ['input_1'], ['lstm_model/Softmax'],
                                   use_cuda=False)

        # validate quantsim
        # get ops and make sure we have a quantized op added to the conditional block
        quantized_graph_op_names = self._get_quant_ops_from_tf_graph(sim.session.graph)

        # _ = tf.summary.FileWriter('./lstm_tm', sess.graph)
        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names, 'lstm_stacked', True, True)
        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names, 'last_lstm')

        # validate forward pass
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim, 9, 14)
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim, 9, 14)

        # close tf sessions
        sess.close()
        sim.session.close()
        del sim
//  [python]
def test_quantize_lstm_deepspeech_time_major_true_quantsim_and_forward_pass(self):
        """ Test connected graph construction on a model with lstm op """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add a LSTM layer with 12 internal units.
            x = tf.keras.layers.LSTM(12,
                                     unit_forget_bias=False,
                                     time_major=True,
                                     return_sequences=True,
                                     name='lstm_stacked')(inputs)
            x2 = tf.keras.layers.LSTM(12, name='last_lstm')(x)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="lstm_model")(x2)

        init = tf.global_variables_initializer()
        sess.run(init)
        # _ = tf.summary.FileWriter('./lstm', sess.graph)

        sim = QuantizationSimModel(sess, ['input_1'], ['lstm_model/Softmax'],
                                   use_cuda=False)

        # validate quantsim
        # get ops and make sure we have a quantized op added to the conditional block
        quantized_graph_op_names = self._get_quant_ops_from_tf_graph(sim.session.graph)

        # _ = tf.summary.FileWriter('./lstm_tm', sess.graph)
        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names, 'lstm_stacked', True, True)
        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names, 'last_lstm')

        # validate forward pass
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim, 9, 14)
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim, 9, 14)

        # close tf sessions
        sess.close()
        sim.session.close()
        del sim
//  [python]
def load_and_save_temporary_tensorflow_model(model_name, model_ckpt, example_img= '../../example/img.jpg'):
    """ Loads and saves a TensorFlow model. """
    image_files = [example_img]
    eval_ckpt_driver = eval_ckpt_main.EvalCkptDriver(model_name)
    with tf.Graph().as_default(), tf.Session() as sess:
        images, labels = eval_ckpt_driver.build_dataset(image_files, [0] * len(image_files), False)
        probs = eval_ckpt_driver.build_model(images, is_training=False)
        sess.run(tf.initialize_all_variables())
        print(model_ckpt)
        eval_ckpt_driver.restore_model(sess, model_ckpt)
        tf.train.Saver().save(sess, 'tmp/model.ckpt')
//  [python]
def load_and_save_temporary_tensorflow_model(model_name, model_ckpt, example_img= '../../example/img.jpg'):
    """ Loads and saves a TensorFlow model. """
    image_files = [example_img]
    eval_ckpt_driver = eval_ckpt_main.EvalCkptDriver(model_name)
    with tf.Graph().as_default(), tf.Session() as sess:
        images, labels = eval_ckpt_driver.build_dataset(image_files, [0] * len(image_files), False)
        probs = eval_ckpt_driver.build_model(images, is_training=False)
        sess.run(tf.global_variables_initializer())
        print(model_ckpt)
        eval_ckpt_driver.restore_model(sess, model_ckpt)
        tf.train.Saver().save(sess, 'tmp/model.ckpt')
//  [python]
def test_quantize_lstm_deepspeech_time_major_false_quantsim_and_forward_pass(self):
        """ Test connected graph construction on a model with lstm op """
        tf.reset_default_graph()

        sess = tf.Session()
        np.random.seed(0)
        tf.set_random_seed(0)

        with sess.graph.as_default():
            inputs = tf.keras.Input(shape=(3, 100))

            # Add a LSTM layer with 12 internal units.
            x, state_h, state_c = tf.keras.layers.LSTM(12, return_state=True,
                                                       return_sequences=True,
                                                       name='lstm_stacked')(inputs)
            x2 = tf.keras.layers.LSTM(12, name='last_lstm')(x)
            _ = tf.keras.layers.Dense(12, activation=tf.nn.softmax,
                                      name="lstm_model")(x2)

        init = tf.initialize_all_variables()
        sess.run(init)

        sim = QuantizationSimModel(sess, ['input_1'], ['lstm_model/Softmax'],
                                   use_cuda=False)

        # validate quantsim
        # get ops and make sure we have a quantized op added to the conditional block
        quantized_graph_op_names = self._get_quant_ops_from_tf_graph(sim.session.graph)

        # _ = tf.summary.FileWriter('./lstm_tm', sess.graph)
        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names, 'lstm_stacked', True, False)
        self.validate_internal_lstm_quantisim_nodes(quantized_graph_op_names, 'last_lstm')

        # validate forward pass
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim, 9, 14)
        self.validate_general_lstm_forward_pass_and_encoding(sess, sim, 9, 14)

        # close tf sessions
        sess.close()
        sim.session.close()
        del sim