jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def load_model(sess, model_path):
    # TODO(global variables ?? how about _adam weights)
    variables = tf.global_variables()
    var_keep_dic = get_variables_in_checkpoint_file(model_path)
    if 'global_step' in var_keep_dic:
        var_keep_dic.pop('global_step')

    # vis_var_keep_dic = []
    variables_to_restore = []
    for v in variables:
        if v.name.split(':')[0] in var_keep_dic:
            # print('Varibles restored: %s' % v.name)
            variables_to_restore.append(v)
            # vis_var_keep_dic.append(v.name.split(':')[0])
        else:
            # print('Unrestored Variables: %s' % v.name)
            pass
    # print('Extra Variables in ckpt', set(var_keep_dic) - set(vis_var_keep_dic))

    if len(variables_to_restore) > 0:
        restorer = tf.train.Saver(variables_to_restore)
        restorer.restore(sess, model_path)
    else:
        print('No variables in {} fits the network'.format(model_path))
//  [python]
weight_path = './data/darknet_weights/yolov3.weights'
save_path = './data/darknet_weights/yolov3.ckpt'
anchors = parse_anchors('./data/yolo_anchors.txt')

model = yolov3(80, anchors)
with tf.Session() as sess:
    inputs = tf.placeholder(tf.float32, [1, img_size, img_size, 3])

    with tf.variable_scope('yolov3'):
        feature_map = model.forward(inputs)

    saver = tf.train.Saver(var_list=tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='yolov3'))

    load_ops = load_weights(tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='yolov3'), weight_path)
    sess.run(load_ops)
    saver.save(sess, save_path=save_path)
    print('TensorFlow model checkpoint has been saved to {}'.format(save_path))
//  [python]
weight_path = './data/darknet_weights/yolov3.weights'
save_path = './data/darknet_weights/yolov3.ckpt'
anchors = parse_anchors('./data/yolo_anchors.txt')

model = yolov3(80, anchors)
with tf.Session() as sess:
    inputs = tf.placeholder(tf.float32, [1, img_size, img_size, 3])

    with tf.variable_scope('yolov3'):
        feature_map = model.forward(inputs)

    saver = tf.train.Saver(var_list=tf.global_variables(scope='yolov3'))

    load_ops = load_weights(tf.global_variables(scope='yolov3'), weight_path)
    sess.run(load_ops)
    saver.save(sess, save_path=save_path)
    print('TensorFlow model checkpoint has been saved to {}'.format(save_path))
//  [python]
def _initialize_variables(graph, fetches, initial_variables):
    with tf.Session(graph=graph) as sess:
        if len(tf.global_variables()) == 0:
            return graph, fetches

        if initial_variables:
            sess.run(tf.global_variables_initializer())
        output_graph_def = tf.graph_util.convert_variables_to_constants(sess,
                            graph.as_graph_def(add_shapes=True),
                            [fetch.op.name for fetch in fetches])

    new_graph = tf.Graph()
    with new_graph.as_default():
        tf.import_graph_def(output_graph_def, name='')
        return new_graph, [new_graph.get_tensor_by_name(fetch.name) for fetch in fetches]
//  [python]
def save_variable_specs(fpath):
    '''Saves information about variables such as
    their name, shape, and total parameter number
    fpath: string. output file path

    Writes
    a text file named fpath.
    '''
    def _get_size(shp):
        '''Gets size of tensor shape
        shp: TensorShape

        Returns
        size
        '''
        size = 1
        for d in range(len(shp)):
            size *=shp[d]
        return size

    params, num_params = [], 0
    for v in tf.global_variables():
        params.append("{}==={}".format(v.name, v.shape))
        num_params += _get_size(v.shape)
    print("num_params: ", num_params)
    with open(fpath, 'w') as fout:
        fout.write("num_params: {}\n".format(num_params))
        fout.write("\n".join(params))
    logging.info("Variables info has been saved.")
//  [python]
def vars(self):
        return [var for var in tf.compat.v1.global_variables() if self.name in var.name]
//  [python]
def vars(self):
        return [var for var in tf.global_variables() if self.name in var.name]
//  [python]
def main():
    """模型训练。"""
    from tensorflow.examples.tutorials.mnist import input_data

    mnist = input_data.read_data_sets("../../data/MNIST_data", one_hot=False)
    print(mnist.test.labels.shape)
    print(mnist.train.labels.shape)

    my_setting = Settings()
    with tf.variable_scope(my_setting.model_name):
        model = Model(my_setting)

    # 模型要保存的变量
    var_list = tf.trainable_variables()
    if model.global_step not in var_list:
        var_list.append(model.global_step)
    # 添加 BN 层的均值和方差
    global_vars = tf.all_variables()
    bn_moving_vars = [v for v in global_vars if 'moving_mean' in v.name]
    bn_moving_vars += [v for v in global_vars if 'moving_variance' in v.name]
    var_list += bn_moving_vars
    # 创建Saver
    saver = tf.train.Saver(var_list=var_list)

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    with tf.Session(config=config) as sess:
        if not os.path.exists(my_setting.ckpt_path + 'checkpoint'):
            print("There is no checkpoit, please check out.")
            exit()
        saver.restore(sess, tf.train.latest_checkpoint(my_setting.ckpt_path))
        tic = time.time()
        n_batch = len(mnist.test.labels) // my_setting.batch_size
        predict_labels = list()
        true_labels = list()
        for step in range(n_batch):
            X_batch, y_batch = mnist.test.next_batch(my_setting.batch_size, shuffle=False)
            X_batch = X_batch.reshape([-1, 28, 28, 1])
            pred_label, test_loss, test_acc = sess.run([model.pred_lables, model.test_loss, model.test_acc],
                                           feed_dict={model.X_inputs: X_batch, model.y_inputs: y_batch})
            predict_labels.append(pred_label)
            true_labels.append(y_batch)
        predict_labels = np.hstack(predict_labels)
        true_labels = np.hstack(true_labels)
        acc = np.sum(predict_labels == true_labels) / len(true_labels)
        print("Test sample number = {}, acc = {:.4f}, pass {:.2f}s".format(len(true_labels), acc, time.time() - tic))
//  [python]
def main():
    """模型训练。"""
    from tensorflow.examples.tutorials.mnist import input_data

    mnist = input_data.read_data_sets("../../data/MNIST_data", one_hot=False)
    print(mnist.test.labels.shape)
    print(mnist.train.labels.shape)

    my_setting = Settings()
    with tf.variable_scope(my_setting.model_name):
        model = Model(my_setting)

    # 模型要保存的变量
    var_list = tf.trainable_variables()
    if model.global_step not in var_list:
        var_list.append(model.global_step)
    # 添加 BN 层的均值和方差
    global_vars = tf.global_variables()
    bn_moving_vars = [v for v in global_vars if 'moving_mean' in v.name]
    bn_moving_vars += [v for v in global_vars if 'moving_variance' in v.name]
    var_list += bn_moving_vars
    # 创建Saver
    saver = tf.train.Saver(var_list=var_list)

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    with tf.Session(config=config) as sess:
        if not os.path.exists(my_setting.ckpt_path + 'checkpoint'):
            print("There is no checkpoit, please check out.")
            exit()
        saver.restore(sess, tf.train.latest_checkpoint(my_setting.ckpt_path))
        tic = time.time()
        n_batch = len(mnist.test.labels) // my_setting.batch_size
        predict_labels = list()
        true_labels = list()
        for step in range(n_batch):
            X_batch, y_batch = mnist.test.next_batch(my_setting.batch_size, shuffle=False)
            X_batch = X_batch.reshape([-1, 28, 28, 1])
            pred_label, test_loss, test_acc = sess.run([model.pred_lables, model.test_loss, model.test_acc],
                                           feed_dict={model.X_inputs: X_batch, model.y_inputs: y_batch})
            predict_labels.append(pred_label)
            true_labels.append(y_batch)
        predict_labels = np.hstack(predict_labels)
        true_labels = np.hstack(true_labels)
        acc = np.sum(predict_labels == true_labels) / len(true_labels)
        print("Test sample number = {}, acc = {:.4f}, pass {:.2f}s".format(len(true_labels), acc, time.time() - tic))
//  [python]
def __init__(self, path, at_epoch, at_every_epoch=0):
        self.path = path + "/"
        self.at_epoch = at_epoch
        self.at_every_epoch = at_every_epoch
        self.saver = tf.train.Saver(tf.all_variables())
//  [python]
def __init__(self, path, at_epoch, at_every_epoch=0):
        self.path = path + "/"
        self.at_epoch = at_epoch
        self.at_every_epoch = at_every_epoch
        self.saver = tf.train.Saver(tf.global_variables())
//  [python]
def restore_from_classification_checkpoint_fn(self, feature_extractor_scope):
    """Returns a map of variables to load from a foreign checkpoint.

    Note that this overrides the default implementation in
    ssd_meta_arch.SSDFeatureExtractor which does not work for PNASNet
    checkpoints.

    Args:
      feature_extractor_scope: A scope name for the first stage feature
        extractor.

    Returns:
      A dict mapping variable names (to load from a checkpoint) to variables in
      the model graph.
    """
    variables_to_restore = {}
    for variable in tf.global_variables():
      if variable.op.name.startswith(feature_extractor_scope):
        var_name = variable.op.name.replace(feature_extractor_scope + '/', '')
        var_name += '/ExponentialMovingAverage'
        variables_to_restore[var_name] = variable
    return variables_to_restore
//  [python]
def initialize_uninitialized(
        self,
    ):
        """Only initializes the variables of a TensorFlow session that were not
        already initialized.
        """
        # List all global variables.
        sess = self.sess
        global_vars = tf.global_variables()

        # Find initialized status for all variables.
        is_var_init = [tf.is_variable_initialized(var) for var in global_vars]
        is_initialized = sess.run(is_var_init)

        # List all variables that were not previously initialized.
        not_initialized_vars = [var for (var, init) in zip(global_vars, is_initialized) if not init]
        for v in not_initialized_vars:
            print("[!] not init: {}".format(v.name))
        # Initialize all uninitialized variables found, if any.
        if len(not_initialized_vars):
            sess.run(tf.variables_initializer(not_initialized_vars))
//  [python]
def init_uninitialized_vars(target_vars: List[tf.Variable] = None) -> None:
    # Initialize all tf.Variables that have not already been initialized
    # Equivalent to the following, but more efficient and does not bloat the tf graph:
    # tf.variables_initializer(tf.report_uninitialized_variables()).run()
    assert_tf_initialized()
    if target_vars is None:
        target_vars = tf.global_variables()

    test_vars = []
    test_ops = []
    with tf.control_dependencies(None): # ignore surrounding control_dependencies
        for var in target_vars:
            assert is_tf_expression(var)

            try:
                tf.get_default_graph().get_tensor_by_name(var.name.replace(":0", "/IsVariableInitialized:0"))
            except KeyError:
                # Op does not exist => variable may be uninitialized
                test_vars.append(var)

                with absolute_name_scope(var.name.split(":")[0]):
                    test_ops.append(tf.is_variable_initialized(var))

    init_vars = [var for var, inited in zip(test_vars, run(test_ops)) if not inited]
    run([var.initializer for var in init_vars])
//  [python]
def save_model(self, model_save_path):
    sess = self.sess
    with self.g.as_default():
      saver = tf.train.Saver(tf.all_variables())
    checkpoint_path = os.path.join(model_save_path, 'vae')
    tf.logging.info('saving model %s.', checkpoint_path)
    saver.save(sess, checkpoint_path, 0)
//  [python]
def save_model(self, model_save_path):
    sess = self.sess
    with self.g.as_default():
      saver = tf.train.Saver(tf.global_variables())
    checkpoint_path = os.path.join(model_save_path, 'vae')
    tf.logging.info('saving model %s.', checkpoint_path)
    saver.save(sess, checkpoint_path, 0)
//  [python]
def load_checkpoint(self, checkpoint_path):
    sess = self.sess
    with self.g.as_default():
      saver = tf.train.Saver(tf.all_variables())
    ckpt = tf.train.get_checkpoint_state(checkpoint_path)
    print('loading model', ckpt.model_checkpoint_path)
    tf.logging.info('Loading model %s.', ckpt.model_checkpoint_path)
    saver.restore(sess, ckpt.model_checkpoint_path)
//  [python]
def load_checkpoint(self, checkpoint_path):
    sess = self.sess
    with self.g.as_default():
      saver = tf.train.Saver(tf.global_variables())
    ckpt = tf.train.get_checkpoint_state(checkpoint_path)
    print('loading model', ckpt.model_checkpoint_path)
    tf.logging.info('Loading model %s.', ckpt.model_checkpoint_path)
    saver.restore(sess, ckpt.model_checkpoint_path)
//  [python]
def init_vars_op(sess):
    variables = tf.global_variables()
    init_flag = sess.run(
        tf.stack([tf.is_variable_initialized(v) for v in variables]))
    uninit_variables = [v for v, f in zip(variables, init_flag) if not f]

    print('Initializing vars:')
    print(PP.pformat([v.name for v in uninit_variables]))

    return tf.variables_initializer(uninit_variables)
//  [python]
def main():

    success = True
    for i in range(10):
        g = np.random.randn(LENGTH, ).astype(np.float32)
        b = np.random.randn(LENGTH, ).astype(np.float32)
        m = np.random.randn(LENGTH, ).astype(np.float32)
        v = np.random.randn(LENGTH, ).astype(np.float32)
        b = b - b.min() + 1.0
        v = v - v.min() + 1.0

        IMG_SIZE = 224
        x = np.random.randn(10, LENGTH, IMG_SIZE, IMG_SIZE).astype(np.float32)

        gl_model = GluonModel()

        # ctx = mx.cpu()
        ctx = mx.gpu(0)
        gl_params = gl_model._collect_params_with_prefix()
        gl_params['bn.gamma']._load_init(mx.nd.array(g, ctx), ctx)
        gl_params['bn.beta']._load_init(mx.nd.array(b, ctx), ctx)
        gl_params['bn.running_mean']._load_init(mx.nd.array(m, ctx), ctx)
        gl_params['bn.running_var']._load_init(mx.nd.array(v, ctx), ctx)
        # gl_model.initialize()

        gl_x = mx.nd.array(x, ctx)
        gl_y = gl_model(gl_x).asnumpy()

        xx = tf.placeholder(
            dtype=tf.float32,
            shape=(None, LENGTH, IMG_SIZE, IMG_SIZE),
            name='xx')
        tf_model = tensorflow_model(xx)
        tf_params = {v.name: v for v in tf.all_variables()}
        with tf.Session() as sess:
            sess.run(tf_params['bn/gamma:0'].assign(g))
            sess.run(tf_params['bn/beta:0'].assign(b))
            sess.run(tf_params['bn/moving_mean:0'].assign(m))
            sess.run(tf_params['bn/moving_variance:0'].assign(v))

            tf_y = sess.run(tf_model, feed_dict={xx: x})
        tf.reset_default_graph()

        diff = np.abs(gl_y - tf_y)
        dist = np.sum(diff)
        if dist > 1e-5:
            success = False
            print("i={}, dist={}".format(i, dist))
            # print(gl_y)
            # print(tf_y)

    if success:
        print("All ok.")
//  [python]
def main():

    success = True
    for i in range(10):
        g = np.random.randn(LENGTH, ).astype(np.float32)
        b = np.random.randn(LENGTH, ).astype(np.float32)
        m = np.random.randn(LENGTH, ).astype(np.float32)
        v = np.random.randn(LENGTH, ).astype(np.float32)
        b = b - b.min() + 1.0
        v = v - v.min() + 1.0

        IMG_SIZE = 224
        x = np.random.randn(10, LENGTH, IMG_SIZE, IMG_SIZE).astype(np.float32)

        gl_model = GluonModel()

        # ctx = mx.cpu()
        ctx = mx.gpu(0)
        gl_params = gl_model._collect_params_with_prefix()
        gl_params['bn.gamma']._load_init(mx.nd.array(g, ctx), ctx)
        gl_params['bn.beta']._load_init(mx.nd.array(b, ctx), ctx)
        gl_params['bn.running_mean']._load_init(mx.nd.array(m, ctx), ctx)
        gl_params['bn.running_var']._load_init(mx.nd.array(v, ctx), ctx)
        # gl_model.initialize()

        gl_x = mx.nd.array(x, ctx)
        gl_y = gl_model(gl_x).asnumpy()

        xx = tf.placeholder(
            dtype=tf.float32,
            shape=(None, LENGTH, IMG_SIZE, IMG_SIZE),
            name='xx')
        tf_model = tensorflow_model(xx)
        tf_params = {v.name: v for v in tf.global_variables()}
        with tf.Session() as sess:
            sess.run(tf_params['bn/gamma:0'].assign(g))
            sess.run(tf_params['bn/beta:0'].assign(b))
            sess.run(tf_params['bn/moving_mean:0'].assign(m))
            sess.run(tf_params['bn/moving_variance:0'].assign(v))

            tf_y = sess.run(tf_model, feed_dict={xx: x})
        tf.reset_default_graph()

        diff = np.abs(gl_y - tf_y)
        dist = np.sum(diff)
        if dist > 1e-5:
            success = False
            print("i={}, dist={}".format(i, dist))
            # print(gl_y)
            # print(tf_y)

    if success:
        print("All ok.")
//  [python]
def print_all_variables(train_only=True):
    """Print all trainable and non-trainable variables
    without tl.layers.initialize_global_variables(sess)

    Parameters
    ----------
    train_only : boolean
        If True, only print the trainable variables, otherwise, print all variables.
    """
    # tvar = tf.trainable_variables() if train_only else tf.all_variables()
    if train_only:
        t_vars = tf.trainable_variables()
        print("  [*] printing trainable variables")
    else:
        try: # TF1.0
            t_vars = tf.global_variables()
        except: # TF0.12
            t_vars = tf.all_variables()
        print("  [*] printing global variables")
    for idx, v in enumerate(t_vars):
        print("  var {:3}: {:15}   {}".format(idx, str(v.get_shape()), v.name))
//  [python]
def _update_bn_accumulators(sess, generated, num_accu_examples):
  """Returns True if the accumlators for batch norm were updated.

  Args:
    sess: `tf.Session` object. Checkpoint should already be loaded.
    generated: Output tensor of the generator.
    num_accu_examples: How many examples should be used to update accumulators.

  Returns:
    True if there were accumlators.
  """
  # Create update ops for batch statistic updates for each batch normalization
  # with accumlators.
  update_accu_switches = [v for v in tf.global_variables()
                          if "accu/update_accus" in v.name]
  logging.info("update_accu_switches: %s", update_accu_switches)
  if not update_accu_switches:
    return False
  sess.run([tf.assign(v, 1) for v in update_accu_switches])
  batch_size = generated.shape[0].value
  num_batches = num_accu_examples // batch_size
  for i in range(num_batches):
    if i % 500 == 0:
      logging.info("Updating BN accumulators %d/%d steps.", i, num_batches)
    sess.run(generated)
  sess.run([tf.assign(v, 0) for v in update_accu_switches])
  logging.info("Done updating BN accumulators.")
  return True
//  [python]
def get_vars(self, name=None, only_trainable=True):
        name = name or self.name
        if only_trainable:
            return [v for v in tf.trainable_variables() if name in v.name]
        else:
            return [v for v in tf.global_variables() if name in v.name]
//  [python]
def testNoBatchNormScaleByDefault(self):
    height, width = 224, 224
    num_classes = 1000
    inputs = tf.placeholder(tf.float32, (1, height, width, 3))
    with slim.arg_scope(inception.inception_v2_arg_scope()):
      inception.inception_v2(inputs, num_classes, is_training=False)

    self.assertEqual(tf.global_variables('.*/BatchNorm/gamma:0$'), [])
//  [python]
def sample(args):
    with open(os.path.join(args.save_dir, 'config.pkl'), 'rb') as f:
        saved_args = cPickle.load(f)
    with open(os.path.join(args.save_dir, 'chars_vocab.pkl'), 'rb') as f:
        chars, vocab = cPickle.load(f)
    #Use most frequent char if no prime is given
    if args.prime == '':
        args.prime = chars[0]
    model = Model(saved_args, training=False)
    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        saver = tf.train.Saver(tf.all_variables())
        ckpt = tf.train.get_checkpoint_state(args.save_dir)
        if ckpt and ckpt.model_checkpoint_path:
            saver.restore(sess, ckpt.model_checkpoint_path)
            data = model.sample(sess, chars, vocab, args.n, args.prime,
                               args.sample).encode('utf-8')
            print(data.decode("utf-8"))
//  [python]
def sample(args):
    with open(os.path.join(args.save_dir, 'config.pkl'), 'rb') as f:
        saved_args = cPickle.load(f)
    with open(os.path.join(args.save_dir, 'chars_vocab.pkl'), 'rb') as f:
        chars, vocab = cPickle.load(f)
    #Use most frequent char if no prime is given
    if args.prime == '':
        args.prime = chars[0]
    model = Model(saved_args, training=False)
    with tf.Session() as sess:
        tf.global_variables_initializer().run()
        saver = tf.train.Saver(tf.global_variables())
        ckpt = tf.train.get_checkpoint_state(args.save_dir)
        if ckpt and ckpt.model_checkpoint_path:
            saver.restore(sess, ckpt.model_checkpoint_path)
            data = model.sample(sess, chars, vocab, args.n, args.prime,
                               args.sample).encode('utf-8')
            print(data.decode("utf-8"))
//  [python]
def getMinigoWeightsV1(model):
    """Load and massage Minigo weights to Leela format.

    This version works on older models (v9 or before)
    But was broken when conv bias was removed in v10
    See: https://github.com/tensorflow/minigo/pull/292 and
         https://github.com/gcp/leela-zero/issues/2020
    """
    sess = tf.Session()
    saver = tf.train.import_meta_graph(model+'.meta')
    saver.restore(sess, model)

    trainable_names = []
    for v in tf.trainable_variables():
        trainable_names.append(v.name)

    weights = []
    for v in tf.all_variables():
        if v.name in trainable_names:
            weights.append(v)
        elif 'batch_normalization' in v.name:
            # Moving mean and variance are not trainable, but are needed for the model
            if 'moving_mean' in v.name or 'moving_variance' in v.name:
                weights.append(v)

    # To match the format of V2
    weights_v2_format = []
    for w in weights:
        nparray = w.eval(session=sess)
        weights_v2_format.append((w.name, nparray))
    return weights_v2_format
//  [python]
def getMinigoWeightsV1(model):
    """Load and massage Minigo weights to Leela format.

    This version works on older models (v9 or before)
    But was broken when conv bias was removed in v10
    See: https://github.com/tensorflow/minigo/pull/292 and
         https://github.com/gcp/leela-zero/issues/2020
    """
    sess = tf.Session()
    saver = tf.train.import_meta_graph(model+'.meta')
    saver.restore(sess, model)

    trainable_names = []
    for v in tf.trainable_variables():
        trainable_names.append(v.name)

    weights = []
    for v in tf.global_variables():
        if v.name in trainable_names:
            weights.append(v)
        elif 'batch_normalization' in v.name:
            # Moving mean and variance are not trainable, but are needed for the model
            if 'moving_mean' in v.name or 'moving_variance' in v.name:
                weights.append(v)

    # To match the format of V2
    weights_v2_format = []
    for w in weights:
        nparray = w.eval(session=sess)
        weights_v2_format.append((w.name, nparray))
    return weights_v2_format
//  [python]
def optimistic_restore(session, save_file, only_load_trainable_vars=False,
                       stat_name_prefix='moving_'):
    """Restore variables of model from save_file.

    Argument trainable_vars is used to determine whether to fetch model
    variables & batch-norm statistics OR whether to fetch ALL variables
    (includes model variables, batch-norm statistics, training variables ~
    such as global step and learning rate decay).

    Args:
        session: tf.Session to use in recovery
        save_file: file to load variables from
        trainable_vars: to recover only variables that are trained (model
            variables and running batch-norm statistics) or not (all variables
            including global step)
    """
    reader = tf.train.NewCheckpointReader(save_file)
    saved_shapes = reader.get_variable_to_shape_map()
    if only_load_trainable_vars:
        var_names = sorted([(var.name, var.name.split(':')[0])
                            for var in tf.trainable_variables()
                            if var.name.split(':')[0] in saved_shapes])
        running_stat_names = sorted([(var.name, var.name.split(':')[0])
                                     for var in tf.all_variables()
                                     if (var.name.split(':')[0]
                                         in saved_shapes and
                                         stat_name_prefix in var.name)])
        var_names += running_stat_names
    else:
        var_names = sorted([(var.name, var.name.split(':')[0])
                            for var in tf.all_variables()
                            if var.name.split(':')[0] in saved_shapes])

    print('Loading vars:')
    print(PP.pformat(var_names))

    restore_vars = []
    name2var = dict(zip(map(lambda x: x.name.split(
        ':')[0], tf.all_variables()), tf.all_variables()))

    with tf.variable_scope('', reuse=True):
        for var_name, saved_var_name in var_names:
            curr_var = name2var[saved_var_name]
            var_shape = curr_var.get_shape().as_list()
            if var_shape == saved_shapes[saved_var_name]:
                restore_vars.append(curr_var)
    saver = tf.train.Saver(restore_vars)
    saver.restore(session, save_file)
//  [python]
def optimistic_restore(session, save_file, only_load_trainable_vars=False,
                       stat_name_prefix='moving_'):
    """Restore variables of model from save_file.

    Argument trainable_vars is used to determine whether to fetch model
    variables & batch-norm statistics OR whether to fetch ALL variables
    (includes model variables, batch-norm statistics, training variables ~
    such as global step and learning rate decay).

    Args:
        session: tf.Session to use in recovery
        save_file: file to load variables from
        trainable_vars: to recover only variables that are trained (model
            variables and running batch-norm statistics) or not (all variables
            including global step)
    """
    reader = tf.train.NewCheckpointReader(save_file)
    saved_shapes = reader.get_variable_to_shape_map()
    if only_load_trainable_vars:
        var_names = sorted([(var.name, var.name.split(':')[0])
                            for var in tf.trainable_variables()
                            if var.name.split(':')[0] in saved_shapes])
        running_stat_names = sorted([(var.name, var.name.split(':')[0])
                                     for var in tf.global_variables()
                                     if (var.name.split(':')[0]
                                         in saved_shapes and
                                         stat_name_prefix in var.name)])
        var_names += running_stat_names
    else:
        var_names = sorted([(var.name, var.name.split(':')[0])
                            for var in tf.global_variables()
                            if var.name.split(':')[0] in saved_shapes])

    print('Loading vars:')
    print(PP.pformat(var_names))

    restore_vars = []
    name2var = dict(zip(map(lambda x: x.name.split(
        ':')[0], tf.global_variables()), tf.global_variables()))

    with tf.variable_scope('', reuse=True):
        for var_name, saved_var_name in var_names:
            curr_var = name2var[saved_var_name]
            var_shape = curr_var.get_shape().as_list()
            if var_shape == saved_shapes[saved_var_name]:
                restore_vars.append(curr_var)
    saver = tf.train.Saver(restore_vars)
    saver.restore(session, save_file)
//  [python]
def testBatchNormScale(self):
    height, width = 224, 224
    num_classes = 1000
    inputs = tf.placeholder(tf.float32, (1, height, width, 3))
    with slim.arg_scope(
        inception.inception_v2_arg_scope(batch_norm_scale=True)):
      inception.inception_v2(inputs, num_classes, is_training=False)

    gamma_names = set(
        v.op.name for v in tf.compat.v1.global_variables('.*/BatchNorm/gamma:0$'))
    self.assertGreater(len(gamma_names), 0)
    for v in tf.compat.v1.global_variables('.*/BatchNorm/moving_mean:0$'):
      self.assertIn(v.op.name[:-len('moving_mean')] + 'gamma', gamma_names)
//  [python]
def testBatchNormScale(self):
    height, width = 224, 224
    num_classes = 1000
    inputs = tf.placeholder(tf.float32, (1, height, width, 3))
    with slim.arg_scope(
        inception.inception_v2_arg_scope(batch_norm_scale=True)):
      inception.inception_v2(inputs, num_classes, is_training=False)

    gamma_names = set(
        v.op.name for v in tf.global_variables('.*/BatchNorm/gamma:0$'))
    self.assertGreater(len(gamma_names), 0)
    for v in tf.global_variables('.*/BatchNorm/moving_mean:0$'):
      self.assertIn(v.op.name[:-len('moving_mean')] + 'gamma', gamma_names)
//  [python]
def gen_poem(begin_word):
    batch_size = 1
    print('## loading corpus from %s' % model_dir)
    poems_vector, word_int_map, vocabularies = process_poems(corpus_file)

    input_data = tf.placeholder(tf.int32, [batch_size, None])

    end_points = rnn_model(model='lstm', input_data=input_data, output_data=None, vocab_size=len(
        vocabularies), rnn_size=128, num_layers=2, batch_size=64, learning_rate=lr)

    saver = tf.train.Saver(tf.all_variables())
    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
    with tf.Session() as sess:
        sess.run(init_op)

        checkpoint = tf.train.latest_checkpoint(model_dir)
        saver.restore(sess, checkpoint)

        x = np.array([list(map(word_int_map.get, start_token))])

        [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']],
                                         feed_dict={input_data: x})
        word = begin_word or to_word(predict, vocabularies)
        poem_ = ''

        i = 0
        while word != end_token:
            poem_ += word
            i += 1
            if i > 24:
                break
            x = np.array([[word_int_map[word]]])
            [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']],
                                             feed_dict={input_data: x, end_points['initial_state']: last_state})
            word = to_word(predict, vocabularies)

        return poem_
//  [python]
def gen_poem(begin_word):
    batch_size = 1
    print('## loading corpus from %s' % model_dir)
    poems_vector, word_int_map, vocabularies = process_poems(corpus_file)

    input_data = tf.placeholder(tf.int32, [batch_size, None])

    end_points = rnn_model(model='lstm', input_data=input_data, output_data=None, vocab_size=len(
        vocabularies), rnn_size=128, num_layers=2, batch_size=64, learning_rate=lr)

    saver = tf.train.Saver(tf.global_variables())
    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
    with tf.Session() as sess:
        sess.run(init_op)

        checkpoint = tf.train.latest_checkpoint(model_dir)
        saver.restore(sess, checkpoint)

        x = np.array([list(map(word_int_map.get, start_token))])

        [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']],
                                         feed_dict={input_data: x})
        word = begin_word or to_word(predict, vocabularies)
        poem_ = ''

        i = 0
        while word != end_token:
            poem_ += word
            i += 1
            if i > 24:
                break
            x = np.array([[word_int_map[word]]])
            [predict, last_state] = sess.run([end_points['prediction'], end_points['last_state']],
                                             feed_dict={input_data: x, end_points['initial_state']: last_state})
            word = to_word(predict, vocabularies)

        return poem_
//  [python]
def initialize_graph(self):
        sess = tf.get_default_session()
        global_vars = tf.global_variables()
        is_not_initialized = sess.run([tf.is_variable_initialized(var) for var in global_vars])
        not_initialized_vars = [v for (v, f) in zip(global_vars, is_not_initialized) if not f]
        print([str(i.name) for i in not_initialized_vars]) # only for testing
        if len(not_initialized_vars):
            sess.run(tf.variables_initializer(not_initialized_vars))
        print("trainable vars: {}".format(tf.trainable_variables(scope=self.scope)))
        saver = tf.train.Saver(tf.trainable_variables(scope=self.scope)) # todo double check
        self.sess = sess
        self.saver = saver
//  [python]
def restore_from_classification_checkpoint_fn(
      self,
      checkpoint_path,
      first_stage_feature_extractor_scope,
      second_stage_feature_extractor_scope):
    """Returns callable for loading a checkpoint into the tensorflow graph.

    Note that this overrides the default implementation in
    faster_rcnn_meta_arch.FasterRCNNFeatureExtractor which does not work for
    InceptionResnetV2 checkpoints.

    TODO: revisit whether it's possible to force the `Repeat` namescope as
    created in `_extract_box_classifier_features` to start counting at 2 (e.g.
    `Repeat_2`) so that the default restore_fn can be used.

    Args:
      checkpoint_path: Path to checkpoint to restore.
      first_stage_feature_extractor_scope: A scope name for the first stage
        feature extractor.
      second_stage_feature_extractor_scope: A scope name for the second stage
        feature extractor.

    Returns:
      a callable which takes a tf.Session as input and loads a checkpoint when
        run.
    """
    variables_to_restore = {}
    for variable in tf.all_variables():
      if variable.op.name.startswith(
          first_stage_feature_extractor_scope):
        var_name = variable.op.name.replace(
            first_stage_feature_extractor_scope + '/', '')
        variables_to_restore[var_name] = variable
      if variable.op.name.startswith(
          second_stage_feature_extractor_scope):
        var_name = variable.op.name.replace(
            second_stage_feature_extractor_scope
            + '/InceptionResnetV2/Repeat', 'InceptionResnetV2/Repeat_2')
        var_name = var_name.replace(
            second_stage_feature_extractor_scope + '/', '')
        variables_to_restore[var_name] = variable
    variables_to_restore = (
        variables_helper.get_variables_available_in_checkpoint(
            variables_to_restore, checkpoint_path))
    saver = tf.train.Saver(variables_to_restore)
    def restore(sess):
      saver.restore(sess, checkpoint_path)
    return restore
//  [python]
def restore_from_classification_checkpoint_fn(
      self,
      checkpoint_path,
      first_stage_feature_extractor_scope,
      second_stage_feature_extractor_scope):
    """Returns callable for loading a checkpoint into the tensorflow graph.

    Note that this overrides the default implementation in
    faster_rcnn_meta_arch.FasterRCNNFeatureExtractor which does not work for
    InceptionResnetV2 checkpoints.

    TODO: revisit whether it's possible to force the `Repeat` namescope as
    created in `_extract_box_classifier_features` to start counting at 2 (e.g.
    `Repeat_2`) so that the default restore_fn can be used.

    Args:
      checkpoint_path: Path to checkpoint to restore.
      first_stage_feature_extractor_scope: A scope name for the first stage
        feature extractor.
      second_stage_feature_extractor_scope: A scope name for the second stage
        feature extractor.

    Returns:
      a callable which takes a tf.Session as input and loads a checkpoint when
        run.
    """
    variables_to_restore = {}
    for variable in tf.global_variables():
      if variable.op.name.startswith(
          first_stage_feature_extractor_scope):
        var_name = variable.op.name.replace(
            first_stage_feature_extractor_scope + '/', '')
        variables_to_restore[var_name] = variable
      if variable.op.name.startswith(
          second_stage_feature_extractor_scope):
        var_name = variable.op.name.replace(
            second_stage_feature_extractor_scope
            + '/InceptionResnetV2/Repeat', 'InceptionResnetV2/Repeat_2')
        var_name = var_name.replace(
            second_stage_feature_extractor_scope + '/', '')
        variables_to_restore[var_name] = variable
    variables_to_restore = (
        variables_helper.get_variables_available_in_checkpoint(
            variables_to_restore, checkpoint_path))
    saver = tf.train.Saver(variables_to_restore)
    def restore(sess):
      saver.restore(sess, checkpoint_path)
    return restore
//  [python]
def begin(self):
    global_variables = tf.global_variables()
    var_names = [re.sub(':[0-9]$', '', var.name) for var in global_variables]
    restore_status = {var_name: False for var_name in var_names}
    for ckpt_path in self._ckpt_path_list:
      logging.info('read variable from %s' % ckpt_path)
      ckpt_reader = tf.train.NewCheckpointReader(ckpt_path)
      ckpt_var2shape_map = ckpt_reader.get_variable_to_shape_map()
      # ckpt_var2shape_map.pop(tf.GraphKeys.GLOBAL_STEP, None)
      name2var = {}
      for var in global_variables:
        var_name = re.sub(':[0-9]$', '', var.name)
        if var_name in ckpt_var2shape_map:
          if restore_status[var_name]:
            logging.warning(
                'variable %s find in more than one checkpoint, skipped %s' %
                (var_name, ckpt_path))
            continue
          name2var[var_name] = var
          restore_status[var_name] = True
      saver = tf.train.Saver(name2var)
      self._saver_list.append(saver)

    restore_check = True
    for var_name, stat in six.iteritems(restore_status):
      if not stat:
        logging.error('var %s not find in checkpoints' % var_name)
        restore_check = False

    assert restore_check, 'failed to find all variables in checkpoints provided'
//  [python]
def initialize(sess=None):
    """Initialize all the uninitialized variables in the global scope."""
    # new_variables = set(tf.global_variables()) - ALREADY_INITIALIZED
    # print(new_variables)
    # get_session().run(tf.variables_initializer(new_variables))
    # ALREADY_INITIALIZED.update(new_variables)
    sess = sess or get_session()
    uninit_vars = []
    for var in tf.all_variables():
        try:
            sess.run(var)
        except tf.errors.FailedPreconditionError:
            # print('uninitialized variables: ', var)
            uninit_vars.append(var)

    if len(uninit_vars) == 0:
        print("None variables need to be initialized...")
    else:
        print('Variables to be initialized: ')
        print(uninit_vars)
    sess.run(tf.variables_initializer(uninit_vars))
//  [python]
def initialize(sess=None):
    """Initialize all the uninitialized variables in the global scope."""
    # new_variables = set(tf.global_variables()) - ALREADY_INITIALIZED
    # print(new_variables)
    # get_session().run(tf.variables_initializer(new_variables))
    # ALREADY_INITIALIZED.update(new_variables)
    sess = sess or get_session()
    uninit_vars = []
    for var in tf.global_variables():
        try:
            sess.run(var)
        except tf.errors.FailedPreconditionError:
            # print('uninitialized variables: ', var)
            uninit_vars.append(var)

    if len(uninit_vars) == 0:
        print("None variables need to be initialized...")
    else:
        print('Variables to be initialized: ')
        print(uninit_vars)
    sess.run(tf.variables_initializer(uninit_vars))
//  [python]
def restore_model(self, sess, ckpt_dir):
    """Restore variables from checkpoint dir."""
    checkpoint = tf.train.latest_checkpoint(ckpt_dir)
    ema = tf.train.ExponentialMovingAverage(decay=0.9999)
    ema_vars = tf.trainable_variables() + tf.get_collection('moving_vars')
    for v in tf.all_variables():
      if 'moving_mean' in v.name or 'moving_variance' in v.name:
        ema_vars.append(v)
    ema_vars = list(set(ema_vars))
    var_dict = ema.variables_to_restore(ema_vars)
    saver = tf.train.Saver(var_dict, max_to_keep=1)
    saver.restore(sess, checkpoint)
//  [python]
def restore_model(self, sess, ckpt_dir):
    """Restore variables from checkpoint dir."""
    checkpoint = tf.train.latest_checkpoint(ckpt_dir)
    ema = tf.train.ExponentialMovingAverage(decay=0.9999)
    ema_vars = tf.trainable_variables() + tf.get_collection('moving_vars')
    for v in tf.global_variables():
      if 'moving_mean' in v.name or 'moving_variance' in v.name:
        ema_vars.append(v)
    ema_vars = list(set(ema_vars))
    var_dict = ema.variables_to_restore(ema_vars)
    saver = tf.train.Saver(var_dict, max_to_keep=1)
    saver.restore(sess, checkpoint)
//  [python]
def print_all_variables(train_only=False):
    """Print all trainable and non-trainable variables
    without tl.layers.initialize_global_variables(sess)

    Parameters
    ----------
    train_only : boolean
        If True, only print the trainable variables, otherwise, print all variables.
    """
    # tvar = tf.trainable_variables() if train_only else tf.all_variables()
    if train_only:
        t_vars = tf.trainable_variables()
        print("  [*] printing trainable variables")
    else:
        try:  # TF1.0+
            t_vars = tf.global_variables()
        except:  # TF0.12
            t_vars = tf.all_variables()
        print("  [*] printing global variables")
    for idx, v in enumerate(t_vars):
        print("  var {:3}: {:15}   {}".format(idx, str(v.get_shape()), v.name))
//  [python]
def retrieve_generator_vars(self):
        all_vars = tf.compat.v1.global_variables()
        generate_vars = [var for var in all_vars if 'embedding' in var.name or "g_" in var.name]
        return generate_vars
//  [python]
def retrieve_generator_vars(self):
        all_vars = tf.global_variables()
        generate_vars = [var for var in all_vars if 'embedding' in var.name or "g_" in var.name]
        return generate_vars
//  [python]
def get_variables_with_name(name, train_only=True, printable=False):
    """Get variable list by a given name scope.

    Examples
    ---------
    >>> dense_vars = tl.layers.get_variable_with_name('dense', True, True)
    """
    print("  [*] geting variables with %s" % name)
    # tvar = tf.trainable_variables() if train_only else tf.all_variables()
    if train_only:
        t_vars = tf.trainable_variables()
    else:
        try:  # TF1.0+
            t_vars = tf.global_variables()
        except:  # TF0.12
            t_vars = tf.all_variables()

    d_vars = [var for var in t_vars if name in var.name]
    if printable:
        for idx, v in enumerate(d_vars):
            print("  got {:3}: {:15}   {}".format(idx, v.name, str(v.get_shape())))
    return d_vars
//  [python]
def run_training():
    if not os.path.exists(FLAGS.model_dir):
        os.makedirs(FLAGS.model_dir)

    poems_vector, word_to_int, vocabularies = process_poems(FLAGS.file_path)
    batches_inputs, batches_outputs = generate_batch(FLAGS.batch_size, poems_vector, word_to_int)

    input_data = tf.placeholder(tf.int32, [FLAGS.batch_size, None])
    output_targets = tf.placeholder(tf.int32, [FLAGS.batch_size, None])

    end_points = rnn_model(model='lstm', input_data=input_data, output_data=output_targets, vocab_size=len(
        vocabularies), rnn_size=128, num_layers=2, batch_size=64, learning_rate=FLAGS.learning_rate)

    saver = tf.train.Saver(tf.all_variables())
    init_op = tf.group(tf.initialize_all_variables(), tf.local_variables_initializer())
    with tf.Session() as sess:
        # sess = tf_debug.LocalCLIDebugWrapperSession(sess=sess)
        # sess.add_tensor_filter("has_inf_or_nan", tf_debug.has_inf_or_nan)
        sess.run(init_op)

        start_epoch = 0
        checkpoint = tf.train.latest_checkpoint(FLAGS.model_dir)
        if checkpoint:
            saver.restore(sess, checkpoint)
            print("## restore from the checkpoint {0}".format(checkpoint))
            start_epoch += int(checkpoint.split('-')[-1])
        print('## start training...')
        try:
            n_chunk = len(poems_vector) // FLAGS.batch_size
            for epoch in range(start_epoch, FLAGS.epochs):
                n = 0
                for batch in range(n_chunk):
                    loss, _, _ = sess.run([
                        end_points['total_loss'],
                        end_points['last_state'],
                        end_points['train_op']
                    ], feed_dict={input_data: batches_inputs[n], output_targets: batches_outputs[n]})
                    n += 1
                    print('Epoch: %d, batch: %d, training loss: %.6f' % (epoch, batch, loss))
                if epoch % 6 == 0:
                    saver.save(sess, os.path.join(FLAGS.model_dir, FLAGS.model_prefix), global_step=epoch)
        except KeyboardInterrupt:
            print('## Interrupt manually, try saving checkpoint for now...')
            saver.save(sess, os.path.join(FLAGS.model_dir, FLAGS.model_prefix), global_step=epoch)
            print('## Last epoch were saved, next time will start from epoch {}.'.format(epoch))
//  [python]
def run_training():
    if not os.path.exists(FLAGS.model_dir):
        os.makedirs(FLAGS.model_dir)

    poems_vector, word_to_int, vocabularies = process_poems(FLAGS.file_path)
    batches_inputs, batches_outputs = generate_batch(FLAGS.batch_size, poems_vector, word_to_int)

    input_data = tf.placeholder(tf.int32, [FLAGS.batch_size, None])
    output_targets = tf.placeholder(tf.int32, [FLAGS.batch_size, None])

    end_points = rnn_model(model='lstm', input_data=input_data, output_data=output_targets, vocab_size=len(
        vocabularies), rnn_size=128, num_layers=2, batch_size=64, learning_rate=FLAGS.learning_rate)

    saver = tf.train.Saver(tf.global_variables())
    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
    with tf.Session() as sess:
        # sess = tf_debug.LocalCLIDebugWrapperSession(sess=sess)
        # sess.add_tensor_filter("has_inf_or_nan", tf_debug.has_inf_or_nan)
        sess.run(init_op)

        start_epoch = 0
        checkpoint = tf.train.latest_checkpoint(FLAGS.model_dir)
        if checkpoint:
            saver.restore(sess, checkpoint)
            print("## restore from the checkpoint {0}".format(checkpoint))
            start_epoch += int(checkpoint.split('-')[-1])
        print('## start training...')
        try:
            n_chunk = len(poems_vector) // FLAGS.batch_size
            for epoch in range(start_epoch, FLAGS.epochs):
                n = 0
                for batch in range(n_chunk):
                    loss, _, _ = sess.run([
                        end_points['total_loss'],
                        end_points['last_state'],
                        end_points['train_op']
                    ], feed_dict={input_data: batches_inputs[n], output_targets: batches_outputs[n]})
                    n += 1
                    print('Epoch: %d, batch: %d, training loss: %.6f' % (epoch, batch, loss))
                if epoch % 6 == 0:
                    saver.save(sess, os.path.join(FLAGS.model_dir, FLAGS.model_prefix), global_step=epoch)
        except KeyboardInterrupt:
            print('## Interrupt manually, try saving checkpoint for now...')
            saver.save(sess, os.path.join(FLAGS.model_dir, FLAGS.model_prefix), global_step=epoch)
            print('## Last epoch were saved, next time will start from epoch {}.'.format(epoch))
//  [python]
def main():
    """Create the model and start the training."""
    args = get_arguments()
    
    # Default image.
    image_batch = tf.constant(0, tf.float32, shape=[1, 321, 321, 3]) 
    # Create network.
    net = DeepLabResNetModel({'data': image_batch})
    var_list = tf.global_variables()
          
    # Set up tf session and initialize variables. 
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    
    with tf.Session(config=config) as sess:
          init = tf.global_variables_initializer()
          sess.run(init)
          
          # Loading .npy weights.
          net.load(args.npy_path, sess)
          
          # Saver for converting the loaded weights into .ckpt.
          saver = tf.train.Saver(var_list=var_list, write_version=1)
          save(saver, sess, args.save_dir)