jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, args):
        self.args = args
        if not args.is_training:
            self.args.batch_size = 1
        if args.hidden_act == 'tanh':
            self.hidden_act = self.tanh
        elif args.hidden_act == 'relu':
            self.hidden_act = self.relu
        else:
            raise NotImplementedError

        self.build_model()
        self.saver = tf.train.Saver(tf.all_variables(), max_to_keep=20)
//  [python]
def __init__(self, args):
        self.args = args
        if not args.is_training:
            self.args.batch_size = 1
        if args.hidden_act == 'tanh':
            self.hidden_act = self.tanh
        elif args.hidden_act == 'relu':
            self.hidden_act = self.relu
        else:
            raise NotImplementedError

        self.build_model()
        self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=20)
//  [python]
def get_state_variables():
    """
    Return set of all tensorflow variables except:

        - Variables in tf.GraphKeys.MODEL_VARIABLES collection
        - Variables in DO_NOT_SAVE collection
    """
    g = tf.get_default_graph()
    vars = set(tf.all_variables())
    vars -= get_model_variables()
    vars -= set(g.get_collection(DO_NOT_SAVE))
    return vars
//  [python]
def get_state_variables():
    """
    Return set of all tensorflow variables except:

        - Variables in tf.GraphKeys.MODEL_VARIABLES collection
        - Variables in DO_NOT_SAVE collection
    """
    g = tf.get_default_graph()
    vars = set(tf.global_variables())
    vars -= get_model_variables()
    vars -= set(g.get_collection(DO_NOT_SAVE))
    return vars
//  [python]
def log_moving_variable():
    def moving_name(name:str):
        name = name[:-2]
        a = name.find("/")
        if a<=0:
            return name+"_moving"
        name_a = name[:a]
        name_b = name[a:]
        return name_a+"_moving"+name_b

    for v in tf.all_variables():
        if "moving_mean" in v.name or "moving_variance" in v.name:
            wsummary.histogram_or_scalar(v,moving_name(v.name))
//  [python]
def log_moving_variable():
    def moving_name(name:str):
        name = name[:-2]
        a = name.find("/")
        if a<=0:
            return name+"_moving"
        name_a = name[:a]
        name_b = name[a:]
        return name_a+"_moving"+name_b

    for v in tf.global_variables():
        if "moving_mean" in v.name or "moving_variance" in v.name:
            wsummary.histogram_or_scalar(v,moving_name(v.name))
//  [python]
def restore_weights_and_initialize(opt):
	var_list = tf.trainable_variables()
	g_list = tf.global_variables()

	# add batch normalization params into trainable variables 
	bn_moving_vars = [g for g in g_list if 'moving_mean' in g.name]
	bn_moving_vars += [g for g in g_list if 'moving_variance' in g.name]
	var_list +=bn_moving_vars

	# create saver to save and restore weights
	resnet_vars = [v for v in var_list if 'resnet_v1_50' in v.name]
	facenet_vars = [v for v in var_list if 'InceptionResnetV1' in v.name]
	saver_resnet = tf.train.Saver(var_list = resnet_vars)
	saver_facenet = tf.train.Saver(var_list = facenet_vars)

	saver = tf.train.Saver(var_list = resnet_vars + [v for v in var_list if 'fc-' in v.name],max_to_keep = 50)

	# create session
	sess = tf.InteractiveSession(config = opt.config)

	# create summary op
	train_writer = tf.summary.FileWriter(opt.train_summary_path, sess.graph)
	val_writer = tf.summary.FileWriter(opt.val_summary_path, sess.graph)

	# initialization
	tf.global_variables_initializer().run()
	tf.local_variables_initializer().run()

	saver_resnet.restore(sess,opt.R_net_weights)
	saver_facenet.restore(sess,opt.Perceptual_net_weights)

	return saver, train_writer,val_writer, sess
//  [python]
def train(self, data=0, steps=-1, dropout=None, display_step=10, test_step=200, batch_size=10,
	          do_resume=False):  # epochs=-1,
		if data: self.data = data
		steps = 9999999 if steps == -1 else steps
		session = self.session
		# with tf.device(_cpu):

		# import tensorflow.contrib.layers as layers
		# t = tf.verify_tensor_all_finite(t, msg)
		tf.add_check_numerics_ops()
		try:
			self.summaries = tf.summary.merge_all()
		except:
			self.summaries = tf.merge_all_summaries()
		try:
			self.summary_writer = tf.summary.FileWriter(current_logdir(), session.graph)  #
		except:
			self.summary_writer = tf.train.SummaryWriter(current_logdir(), session.graph)  #
		if not dropout: dropout = 1.  # keep all
		x = self.x
		y = self.y
		keep_prob = self.keep_prob
		try:
			saver = tf.train.Saver(tf.all_variables())
		except:
			saver = tf.train.Saver(tf.all_variables())
		snapshot = self.name + str(get_last_tensorboard_run_nr())
		checkpoint = tf.train.latest_checkpoint(checkpoint_dir)
		if do_resume and checkpoint:
			print("LOADING " + checkpoint + " !!!")
			saver.restore(session, checkpoint)
		try:
			session.run([tf.global_variables_initializer()])
		except:
			session.run([tf.initialize_all_variables()])
		step = 0  # show first
		while step < steps:
			batch_xs, batch_ys = self.next_batch(batch_size, session)
			# print("step %d \r" % step)# end=' ')

			# tf.train.shuffle_batch_join(example_list, batch_size, capacity=min_queue_size + batch_size * 16, min_queue_size)
			# Fit training using batch data
			feed_dict = {x: batch_xs, y: batch_ys, keep_prob: dropout, self.train_phase: True}
			loss, _ = session.run([self.cost, self.optimizer], feed_dict=feed_dict)
			if step % display_step == 0:
				seconds = int(time.time()) - start
				# Calculate batch accuracy, loss
				feed = {x: batch_xs, y: batch_ys, keep_prob: 1., self.train_phase: False}
				acc, summary = session.run([self.accuracy, self.summaries], feed_dict=feed)
				# self.summary_writer.add_summary(summary, step) # only test summaries for smoother curve
				print("\rStep {:d} Loss= {:.6f} Accuracy= {:.3f} Time= {:d}s".format(step, loss, acc, seconds), end=' ')
				if str(loss) == "nan": return print("\nLoss gradiant explosion, exiting!!!")  # restore!
			if step % test_step == 0: self.test(step)
			if step % save_step == 0 and step > 0:
				print("SAVING snapshot %s" % snapshot)
				saver.save(session, checkpoint_dir + snapshot + ".ckpt", self.global_step)

			step += 1
		print("\nOptimization Finished!")
		self.test(step, number=10000)
//  [python]
def train(self, data=0, steps=-1, dropout=None, display_step=10, test_step=200, batch_size=10,
	          do_resume=False):  # epochs=-1,
		if data: self.data = data
		steps = 9999999 if steps == -1 else steps
		session = self.session
		# with tf.device(_cpu):

		# import tensorflow.contrib.layers as layers
		# t = tf.verify_tensor_all_finite(t, msg)
		tf.add_check_numerics_ops()
		try:
			self.summaries = tf.summary.merge_all()
		except:
			self.summaries = tf.merge_all_summaries()
		try:
			self.summary_writer = tf.summary.FileWriter(current_logdir(), session.graph)  #
		except:
			self.summary_writer = tf.train.SummaryWriter(current_logdir(), session.graph)  #
		if not dropout: dropout = 1.  # keep all
		x = self.x
		y = self.y
		keep_prob = self.keep_prob
		try:
			saver = tf.train.Saver(tf.global_variables())
		except:
			saver = tf.train.Saver(tf.all_variables())
		snapshot = self.name + str(get_last_tensorboard_run_nr())
		checkpoint = tf.train.latest_checkpoint(checkpoint_dir)
		if do_resume and checkpoint:
			print("LOADING " + checkpoint + " !!!")
			saver.restore(session, checkpoint)
		try:
			session.run([tf.global_variables_initializer()])
		except:
			session.run([tf.initialize_all_variables()])
		step = 0  # show first
		while step < steps:
			batch_xs, batch_ys = self.next_batch(batch_size, session)
			# print("step %d \r" % step)# end=' ')

			# tf.train.shuffle_batch_join(example_list, batch_size, capacity=min_queue_size + batch_size * 16, min_queue_size)
			# Fit training using batch data
			feed_dict = {x: batch_xs, y: batch_ys, keep_prob: dropout, self.train_phase: True}
			loss, _ = session.run([self.cost, self.optimizer], feed_dict=feed_dict)
			if step % display_step == 0:
				seconds = int(time.time()) - start
				# Calculate batch accuracy, loss
				feed = {x: batch_xs, y: batch_ys, keep_prob: 1., self.train_phase: False}
				acc, summary = session.run([self.accuracy, self.summaries], feed_dict=feed)
				# self.summary_writer.add_summary(summary, step) # only test summaries for smoother curve
				print("\rStep {:d} Loss= {:.6f} Accuracy= {:.3f} Time= {:d}s".format(step, loss, acc, seconds), end=' ')
				if str(loss) == "nan": return print("\nLoss gradiant explosion, exiting!!!")  # restore!
			if step % test_step == 0: self.test(step)
			if step % save_step == 0 and step > 0:
				print("SAVING snapshot %s" % snapshot)
				saver.save(session, checkpoint_dir + snapshot + ".ckpt", self.global_step)

			step += 1
		print("\nOptimization Finished!")
		self.test(step, number=10000)
//  [python]
def initialize_uninitialized_variables():
    with tf.name_scope("initialize"):
        uninitialized_names = set(tf.get_default_session().run(tf.report_uninitialized_variables()))
        uninitialized_vars = []
        for var in tf.all_variables():
            if var.name[:-2].encode() in uninitialized_names:
                uninitialized_vars.append(var)

        tf.variables_initializer(uninitialized_vars).run()
//  [python]
def initialize_uninitialized_variables():
    with tf.name_scope("initialize"):
        uninitialized_names = set(tf.get_default_session().run(tf.report_uninitialized_variables()))
        uninitialized_vars = []
        for var in tf.global_variables():
            if var.name[:-2].encode() in uninitialized_names:
                uninitialized_vars.append(var)

        tf.variables_initializer(uninitialized_vars).run()
//  [python]
def main():
    """Create the model and start the evaluation process."""
    args = get_arguments()
    
    # Create queue coordinator.
    coord = tf.train.Coordinator()
    
    # Load reader.
    with tf.name_scope("create_inputs"):
        reader = ImageReader(
            args.data_dir,
            args.data_list,
            None, # No defined input size.
            False, # No random scale.
            False, # No random mirror.
            args.ignore_label,
            IMG_MEAN,
            coord)
        image, label = reader.image, reader.label
    image_batch, label_batch = tf.expand_dims(image, dim=0), tf.expand_dims(label, dim=0) # Add one batch dimension.

    # Create network.
    net = DeepLabResNetModel({'data': image_batch}, is_training=False, num_classes=args.num_classes)

    # Which variables to load.
    restore_var = tf.all_variables()
    
    # Predictions.
    raw_output = net.layers['fc1_voc12']
    raw_output = tf.image.resize_bilinear(raw_output, tf.shape(image_batch)[1:3,])
    raw_output = tf.argmax(raw_output, dimension=3)
    pred = tf.expand_dims(raw_output, dim=3) # Create 4-d tensor.
    
    # mIoU
    pred = tf.reshape(pred, [-1,])
    gt = tf.reshape(label_batch, [-1,])
    weights = tf.cast(tf.less_equal(gt, args.num_classes - 1), tf.int32) # Ignoring all labels greater than or equal to n_classes.
    mIoU, update_op = tf.contrib.metrics.streaming_mean_iou(pred, gt, num_classes=args.num_classes, weights=weights)
    
    # Set up tf session and initialize variables. 
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    init = tf.global_variables_initializer()
    
    sess.run(init)
    sess.run(tf.local_variables_initializer())
    
    # Load weights.
    loader = tf.train.Saver(var_list=restore_var)
    if args.restore_from is not None:
        load(loader, sess, args.restore_from)
    
    # Start queue threads.
    threads = tf.train.start_queue_runners(coord=coord, sess=sess)
    
    # Iterate over training steps.
    for step in range(args.num_steps):
        preds, _ = sess.run([pred, update_op])
        if step % 100 == 0:
            print('step {:d}'.format(step))
    print('Mean IoU: {:.3f}'.format(mIoU.eval(session=sess)))
    coord.request_stop()
    coord.join(threads)
//  [python]
def main():
    """Create the model and start the evaluation process."""
    args = get_arguments()
    
    # Create queue coordinator.
    coord = tf.train.Coordinator()
    
    # Load reader.
    with tf.name_scope("create_inputs"):
        reader = ImageReader(
            args.data_dir,
            args.data_list,
            None, # No defined input size.
            False, # No random scale.
            False, # No random mirror.
            args.ignore_label,
            IMG_MEAN,
            coord)
        image, label = reader.image, reader.label
    image_batch, label_batch = tf.expand_dims(image, dim=0), tf.expand_dims(label, dim=0) # Add one batch dimension.

    # Create network.
    net = DeepLabResNetModel({'data': image_batch}, is_training=False, num_classes=args.num_classes)

    # Which variables to load.
    restore_var = tf.global_variables()
    
    # Predictions.
    raw_output = net.layers['fc1_voc12']
    raw_output = tf.image.resize_bilinear(raw_output, tf.shape(image_batch)[1:3,])
    raw_output = tf.argmax(raw_output, dimension=3)
    pred = tf.expand_dims(raw_output, dim=3) # Create 4-d tensor.
    
    # mIoU
    pred = tf.reshape(pred, [-1,])
    gt = tf.reshape(label_batch, [-1,])
    weights = tf.cast(tf.less_equal(gt, args.num_classes - 1), tf.int32) # Ignoring all labels greater than or equal to n_classes.
    mIoU, update_op = tf.contrib.metrics.streaming_mean_iou(pred, gt, num_classes=args.num_classes, weights=weights)
    
    # Set up tf session and initialize variables. 
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    init = tf.global_variables_initializer()
    
    sess.run(init)
    sess.run(tf.local_variables_initializer())
    
    # Load weights.
    loader = tf.train.Saver(var_list=restore_var)
    if args.restore_from is not None:
        load(loader, sess, args.restore_from)
    
    # Start queue threads.
    threads = tf.train.start_queue_runners(coord=coord, sess=sess)
    
    # Iterate over training steps.
    for step in range(args.num_steps):
        preds, _ = sess.run([pred, update_op])
        if step % 100 == 0:
            print('step {:d}'.format(step))
    print('Mean IoU: {:.3f}'.format(mIoU.eval(session=sess)))
    coord.request_stop()
    coord.join(threads)
//  [python]
def _countInceptionParameters(self):
    """Counts the number of parameters in the inception model at top scope."""
    counter = {}
    for v in tf.global_variables():
      name_tokens = v.op.name.split("/")
      if name_tokens[0] == "InceptionV3":
        name = "InceptionV3/" + name_tokens[1]
        num_params = v.get_shape().num_elements()
        assert num_params
        counter[name] = counter.get(name, 0) + num_params
    return counter
//  [python]
def optimistic_restore(save_file, graph=tf.get_default_graph()):
    print("optimistic_restore")
    reader = tf.train.NewCheckpointReader(save_file)
    saved_shapes = reader.get_variable_to_shape_map()
    var_names = sorted([(var.name, var.name.split(':')[0]) for var in tf.all_variables()
            if var.name.split(':')[0] in saved_shapes])    
    restore_vars = []    
    for var_name, saved_var_name in var_names:            
        curr_var = graph.get_tensor_by_name(var_name)
        var_shape = curr_var.get_shape().as_list()
        if var_shape == saved_shapes[saved_var_name]:
            restore_vars.append(curr_var)
    opt_saver = tf.train.Saver(restore_vars)
    return opt_saver
//  [python]
def optimistic_restore(save_file, graph=tf.get_default_graph()):
    print("optimistic_restore")
    reader = tf.train.NewCheckpointReader(save_file)
    saved_shapes = reader.get_variable_to_shape_map()
    var_names = sorted([(var.name, var.name.split(':')[0]) for var in tf.global_variables()
            if var.name.split(':')[0] in saved_shapes])    
    restore_vars = []    
    for var_name, saved_var_name in var_names:            
        curr_var = graph.get_tensor_by_name(var_name)
        var_shape = curr_var.get_shape().as_list()
        if var_shape == saved_shapes[saved_var_name]:
            restore_vars.append(curr_var)
    opt_saver = tf.train.Saver(restore_vars)
    return opt_saver
//  [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, 'words_vocab.pkl'), 'rb') as f:
        words, vocab = cPickle.load(f)
    model = Model(saved_args, True)
    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)
            for _ in range(args.count):
              print(model.sample(sess, words, vocab, args.n, args.prime, args.sample, args.pick, args.width, args.quiet))
//  [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, 'words_vocab.pkl'), 'rb') as f:
        words, vocab = cPickle.load(f)
    model = Model(saved_args, True)
    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)
            for _ in range(args.count):
              print(model.sample(sess, words, vocab, args.n, args.prime, args.sample, args.pick, args.width, args.quiet))
//  [python]
def main():

    success = True
    for i in range(10):
        # w = np.random.randint(10, size=(64, 3, 7, 7)).astype(np.float32)
        # x = np.random.randint(10, size=(1, 3, 224, 224)).astype(np.float32)
        w = np.random.randn(64, 3, 7, 7).astype(np.float32)
        b = np.random.randn(64, ).astype(np.float32)
        x = np.random.randn(10, 3, 224, 224).astype(np.float32)

        gl_model = GluonModel()

        # ctx = mx.cpu()
        ctx = mx.gpu(0)
        gl_params = gl_model._collect_params_with_prefix()
        gl_params['conv.weight']._load_init(mx.nd.array(w, ctx), ctx)
        gl_params['conv.bias']._load_init(mx.nd.array(b, ctx), ctx)

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

        xx = tf.placeholder(
            dtype=tf.float32,
            shape=(None, 3, 224, 224),
            name='xx')
        tf_model = tensorflow_model(xx)
        tf_params = {v.name: v for v in tf.all_variables()}
        with tf.Session() as sess:
            tf_w = np.transpose(w, axes=(2, 3, 1, 0))
            sess.run(tf_params['conv/kernel:0'].assign(tf_w))
            sess.run(tf_params['conv/bias:0'].assign(b))

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

        dist = np.sum(np.abs(gl_y - tf_y))
        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):
        # w = np.random.randint(10, size=(64, 3, 7, 7)).astype(np.float32)
        # x = np.random.randint(10, size=(1, 3, 224, 224)).astype(np.float32)
        w = np.random.randn(64, 3, 7, 7).astype(np.float32)
        b = np.random.randn(64, ).astype(np.float32)
        x = np.random.randn(10, 3, 224, 224).astype(np.float32)

        gl_model = GluonModel()

        # ctx = mx.cpu()
        ctx = mx.gpu(0)
        gl_params = gl_model._collect_params_with_prefix()
        gl_params['conv.weight']._load_init(mx.nd.array(w, ctx), ctx)
        gl_params['conv.bias']._load_init(mx.nd.array(b, ctx), ctx)

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

        xx = tf.placeholder(
            dtype=tf.float32,
            shape=(None, 3, 224, 224),
            name='xx')
        tf_model = tensorflow_model(xx)
        tf_params = {v.name: v for v in tf.global_variables()}
        with tf.Session() as sess:
            tf_w = np.transpose(w, axes=(2, 3, 1, 0))
            sess.run(tf_params['conv/kernel:0'].assign(tf_w))
            sess.run(tf_params['conv/bias:0'].assign(b))

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

        dist = np.sum(np.abs(gl_y - tf_y))
        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 get_variable_by_name(name, graph=None):
    if len(tf.global_variables()) == 0:
        name = name + ":0"
        return graph.get_tensor_by_name(name)
    else:
        res = None
        all_variable_names = [v.op.name for v in tf.global_variables()]
        if name in all_variable_names:
            res = [v for v in tf.global_variables() if v.op.name == name][0]
        return res
//  [python]
def main():

    success = True
    for i in range(10):
        # gl_w = np.random.randn(32, 1, 7, 7).astype(np.float32)
        tf_w = np.random.randn(7, 7, 32, 1).astype(np.float32)
        x = np.random.randn(10, 32, 224, 224).astype(np.float32)

        gl_model = GluonModel()

        # ctx = mx.cpu()
        ctx = mx.gpu(0)
        gl_params = gl_model._collect_params_with_prefix()
        gl_w = np.transpose(tf_w, axes=(2, 3, 0, 1))
        gl_params['dw_conv.weight']._load_init(mx.nd.array(gl_w, ctx), ctx)

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

        xx = tf.placeholder(
            dtype=tf.float32,
            shape=(None, 32, 224, 224),
            name='xx')
        tf_model = tensorflow_model(xx)
        tf_params = {v.name: v for v in tf.all_variables()}
        with tf.Session() as sess:
            # tf_w = np.transpose(gl_w, axes=(2, 3, 0, 1))
            sess.run(tf_params['dw_conv/dw_kernel:0'].assign(tf_w))

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

        dist = np.sum(np.abs(gl_y - tf_y))
        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):
        # gl_w = np.random.randn(32, 1, 7, 7).astype(np.float32)
        tf_w = np.random.randn(7, 7, 32, 1).astype(np.float32)
        x = np.random.randn(10, 32, 224, 224).astype(np.float32)

        gl_model = GluonModel()

        # ctx = mx.cpu()
        ctx = mx.gpu(0)
        gl_params = gl_model._collect_params_with_prefix()
        gl_w = np.transpose(tf_w, axes=(2, 3, 0, 1))
        gl_params['dw_conv.weight']._load_init(mx.nd.array(gl_w, ctx), ctx)

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

        xx = tf.placeholder(
            dtype=tf.float32,
            shape=(None, 32, 224, 224),
            name='xx')
        tf_model = tensorflow_model(xx)
        tf_params = {v.name: v for v in tf.global_variables()}
        with tf.Session() as sess:
            # tf_w = np.transpose(gl_w, axes=(2, 3, 0, 1))
            sess.run(tf_params['dw_conv/dw_kernel:0'].assign(tf_w))

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

        dist = np.sum(np.abs(gl_y - tf_y))
        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 restore_weights_and_initialize(train_stage_args):
    var_list = tf.trainable_variables()
    g_list = tf.global_variables()

    # add batch normalization params into trainable variables 
    bn_moving_vars = [g for g in g_list if 'moving_mean' in g.name]
    bn_moving_vars += [g for g in g_list if 'moving_variance' in g.name]
    var_list +=bn_moving_vars

    var_id_list = [v for v in var_list if 'id' in v.name and 'stage1' in v.name]
    var_exp_list = [v for v in var_list if 'exp' in v.name and 'stage1' in v.name]
    var_gamma_list = [v for v in var_list if 'gamma' in v.name and 'stage1' in v.name]
    var_rot_list = [v for v in var_list if 'rot' in v.name and 'stage1' in v.name]

    resnet_vars = [v for v in var_list if 'resnet_v1_50' in v.name]
    res_fc = [v for v in var_list if 'fc-id' in v.name or 'fc-ex' in v.name or 'fc-tex' in v.name or 'fc-angles' in v.name or 'fc-gamma' in v.name or 'fc-XY' in v.name or 'fc-Z' in v.name]
    resnet_vars += res_fc

    facerec_vars = [v for v in var_list if 'InceptionResnetV1' in v.name]
    

    saver_resnet = tf.train.Saver(var_list = resnet_vars)
    saver_facerec = tf.train.Saver(var_list = facerec_vars)
    saver_id = tf.train.Saver(var_list = var_id_list,max_to_keep = 100)
    saver_exp = tf.train.Saver(var_list = var_exp_list,max_to_keep = 100)
    saver_gamma = tf.train.Saver(var_list = var_gamma_list,max_to_keep = 100)
    saver_rot = tf.train.Saver(var_list = var_rot_list,max_to_keep = 100)
    

    saver_resnet.restore(tf.get_default_session(),os.path.join('./training/pretrained_weights/recon_net','FaceReconModel'))
    saver_facerec.restore(tf.get_default_session(),'./training/pretrained_weights/id_net/model-20170512-110547.ckpt-250000')
    saver_id.restore(tf.get_default_session(),'./vae/weights/id/stage1_epoch_395.ckpt')
    saver_exp.restore(tf.get_default_session(),'./vae/weights/exp/stage1_epoch_395.ckpt')
    saver_gamma.restore(tf.get_default_session(),'./vae/weights/gamma/stage1_epoch_395.ckpt')
    saver_rot.restore(tf.get_default_session(),'./vae/weights/rot/stage1_epoch_395.ckpt')

    if train_stage_args.func_name == 'training.training_utils.training_stage2':
        parser_vars = [v for v in var_list if 'FaceParser' in v.name]
        saver_parser = tf.train.Saver(var_list = parser_vars)
        saver_parser.restore(tf.get_default_session(),os.path.join('./training/pretrained_weights/parsing_net','faceparser_public'))
//  [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]
def testDistributedReplicatedSavableVars(self):
    test_util.monkey_patch_base_cluster_manager()
    params = benchmark_cnn.make_params(
        variable_update='distributed_replicated',
        model='inception4',
        data_name='imagenet',
        data_dir=os.path.join(platforms_util.get_test_data_dir(),
                              'fake_tf_record_data'),
        job_name='worker',
        worker_hosts='w1,w2,w3,w4',
        ps_hosts='p1',
        datasets_use_prefetch=False)

    bench = benchmark_cnn.BenchmarkCNN(params)
    with tf.Graph().as_default():
      bench._build_model()
      savable_vars = bench.variable_mgr.savable_variables()
      # Assert all global variables are in savable_vars
      for v in tf.global_variables():
        if not v.name.startswith(
            variable_mgr_util.PS_SHADOW_VAR_PREFIX + '/v0'):
          self.assertEqual(v.name, 'global_step:0')
        name = bench.variable_mgr._strip_port(v.name)
        if name.startswith(variable_mgr_util.PS_SHADOW_VAR_PREFIX):
          name = name[len(variable_mgr_util.PS_SHADOW_VAR_PREFIX + '/'):]
        self.assertIn(name, savable_vars)
        self.assertIn(savable_vars[name], tf.global_variables())
      # Assert all local variables on the first tower are in savable_vars
      for v in tf.local_variables():
        if v.name.startswith('v0/'):
          name = bench.variable_mgr._strip_port(v.name)
          self.assertIn(name, savable_vars)
//  [python]
def _build(self, inputs):
    (shared_inputs, extra_policy_inputs) = inputs
    policy_in = tf.concat([shared_inputs, extra_policy_inputs], axis=1)

    policy = snt.nets.MLP(
        output_sizes=self._policy_layers,
        activation=self._activation,
        name='policy_mlp')(
            policy_in)

    # Sample an action from the policy logits.
    action = tf.multinomial(policy, num_samples=1, output_dtype=tf.int32)
    action = tf.squeeze(action, 1)  # [B, 1] -> [B]

    if self._policy_clip_abs_value > 0:
      policy = snt.clip_gradient(
          net=policy,
          clip_value_min=-self._policy_clip_abs_value,
          clip_value_max=self._policy_clip_abs_value)

    baseline_in = tf.concat([shared_inputs, tf.stop_gradient(policy)], axis=1)
    baseline = snt.nets.MLP(
        self._baseline_layers,
        activation=self._activation,
        name='baseline_mlp')(
            baseline_in)
    baseline = tf.squeeze(baseline, axis=-1)  # [B, 1] -> [B]

    if self._policy_clip_abs_value > 0:
      baseline = snt.clip_gradient(
          net=baseline,
          clip_value_min=-self._policy_clip_abs_value,
          clip_value_max=self._policy_clip_abs_value)

    outputs = PolicyOutputs(
        policy=policy,
        action=action,
        baseline=baseline)

    return outputs
//  [python]
def _build(self, inputs):
    (shared_inputs, extra_policy_inputs) = inputs
    policy_in = tf.concat([shared_inputs, extra_policy_inputs], axis=1)

    policy = snt.nets.MLP(
        output_sizes=self._policy_layers,
        activation=self._activation,
        name='policy_mlp')(
            policy_in)

    # Sample an action from the policy logits.
    action = tf.random.categorical(policy, num_samples=1)
    action = tf.squeeze(action, 1)  # [B, 1] -> [B]

    if self._policy_clip_abs_value > 0:
      policy = snt.clip_gradient(
          net=policy,
          clip_value_min=-self._policy_clip_abs_value,
          clip_value_max=self._policy_clip_abs_value)

    baseline_in = tf.concat([shared_inputs, tf.stop_gradient(policy)], axis=1)
    baseline = snt.nets.MLP(
        self._baseline_layers,
        activation=self._activation,
        name='baseline_mlp')(
            baseline_in)
    baseline = tf.squeeze(baseline, axis=-1)  # [B, 1] -> [B]

    if self._policy_clip_abs_value > 0:
      baseline = snt.clip_gradient(
          net=baseline,
          clip_value_min=-self._policy_clip_abs_value,
          clip_value_max=self._policy_clip_abs_value)

    outputs = PolicyOutputs(
        policy=policy,
        action=action,
        baseline=baseline)

    return outputs
//  [python]
def evaluate():
  tf.set_random_seed(0)  # for reproducibility
  with tf.Graph().as_default():
    # convert feature_names and feature_sizes to lists of values
    feature_names, feature_sizes = utils.GetListOfFeatureNamesAndSizes(
        FLAGS.feature_names, FLAGS.feature_sizes)

    if FLAGS.frame_features:
      reader = readers.YT8MFrameFeatureReader(feature_names=feature_names,
                                              feature_sizes=feature_sizes)
    else:
      reader = readers.YT8MAggregatedFeatureReader(feature_names=feature_names,
                                                   feature_sizes=feature_sizes)

    model = find_class_by_name(FLAGS.model,
        [frame_level_models, video_level_models])()
    label_loss_fn = find_class_by_name(FLAGS.label_loss, [losses])()

    if FLAGS.eval_data_pattern is "":
      raise IOError("'eval_data_pattern' was not specified. " +
                     "Nothing to evaluate.")

    build_graph(
        reader=reader,
        model=model,
        eval_data_pattern=FLAGS.eval_data_pattern,
        label_loss_fn=label_loss_fn,
        num_readers=FLAGS.num_readers,
        batch_size=FLAGS.batch_size)
    logging.info("built evaluation graph")
    video_id_batch = tf.get_collection("video_id_batch")[0]
    prediction_batch = tf.get_collection("predictions")[0]
    label_batch = tf.get_collection("labels")[0]
    loss = tf.get_collection("loss")[0]
    summary_op = tf.get_collection("summary_op")[0]

    saver = tf.train.Saver(tf.all_variables())
    summary_writer = tf.summary.FileWriter(
        FLAGS.train_dir, graph=tf.get_default_graph())

    evl_metrics = eval_util.EvaluationMetrics(reader.num_classes, FLAGS.top_k)

    last_global_step_val = -1
    while True:
      last_global_step_val = evaluation_loop(video_id_batch, prediction_batch,
                                             label_batch, loss, summary_op,
                                             saver, summary_writer, evl_metrics,
                                             last_global_step_val)
      if FLAGS.run_once:
        break
//  [python]
def evaluate():
  tf.set_random_seed(0)  # for reproducibility
  with tf.Graph().as_default():
    # convert feature_names and feature_sizes to lists of values
    feature_names, feature_sizes = utils.GetListOfFeatureNamesAndSizes(
        FLAGS.feature_names, FLAGS.feature_sizes)

    if FLAGS.frame_features:
      reader = readers.YT8MFrameFeatureReader(feature_names=feature_names,
                                              feature_sizes=feature_sizes)
    else:
      reader = readers.YT8MAggregatedFeatureReader(feature_names=feature_names,
                                                   feature_sizes=feature_sizes)

    model = find_class_by_name(FLAGS.model,
        [frame_level_models, video_level_models])()
    label_loss_fn = find_class_by_name(FLAGS.label_loss, [losses])()

    if FLAGS.eval_data_pattern is "":
      raise IOError("'eval_data_pattern' was not specified. " +
                     "Nothing to evaluate.")

    build_graph(
        reader=reader,
        model=model,
        eval_data_pattern=FLAGS.eval_data_pattern,
        label_loss_fn=label_loss_fn,
        num_readers=FLAGS.num_readers,
        batch_size=FLAGS.batch_size)
    logging.info("built evaluation graph")
    video_id_batch = tf.get_collection("video_id_batch")[0]
    prediction_batch = tf.get_collection("predictions")[0]
    label_batch = tf.get_collection("labels")[0]
    loss = tf.get_collection("loss")[0]
    summary_op = tf.get_collection("summary_op")[0]

    saver = tf.train.Saver(tf.global_variables())
    summary_writer = tf.summary.FileWriter(
        FLAGS.train_dir, graph=tf.get_default_graph())

    evl_metrics = eval_util.EvaluationMetrics(reader.num_classes, FLAGS.top_k)

    last_global_step_val = -1
    while True:
      last_global_step_val = evaluation_loop(video_id_batch, prediction_batch,
                                             label_batch, loss, summary_op,
                                             saver, summary_writer, evl_metrics,
                                             last_global_step_val)
      if FLAGS.run_once:
        break
//  [python]
def sample_sequence(*, hparams, length, start_token=None, batch_size=None, context=None, temperature=1, top_k=0):
    if start_token is None:
        assert context is not None, 'Specify exactly one of start_token and context!'
    else:
        assert context is None, 'Specify exactly one of start_token and context!'
        context = tf.fill([batch_size, 1], start_token)

    # get length of prefix
    prefix_len = tf.shape(context)[-1]

    def step(hparams, tokens, past=None):
        lm_output = model.model(
            hparams=hparams, X=tokens, past=past, reuse=tf.AUTO_REUSE)

        logits = lm_output['logits'][:, :, :hparams.n_vocab]
        presents = lm_output['present']
        presents.set_shape(model.past_shape(
            hparams=hparams, batch_size=batch_size))
        return {
            'logits': logits,
            'presents': presents,
        }

    with tf.name_scope('sample_sequence'):
        # Don't feed the last context token -- leave that to the loop below
        # TODO: Would be slightly faster if we called step on the entire context,
        # rather than leaving the last token transformer calculation to the while loop.
        context_output = step(hparams, context[:, :-1])

        def body(past, prev, output):
            next_outputs = step(hparams, prev[:, tf.newaxis], past=past)
            logits = next_outputs['logits'][:, -1, :] / \
                tf.to_float(temperature)
            logits = top_k_logits(logits, k=top_k)
            samples = tf.multinomial(
                logits, num_samples=1, output_dtype=tf.int32)
            return [
                tf.concat([past, next_outputs['presents']], axis=-2),
                tf.squeeze(samples, axis=[1]),
                tf.concat([output, samples], axis=1),
            ]

        def cond(*args):
            return True
        with tf.control_dependencies([prefix_len]):
            _, _, tokens = tf.while_loop(
                cond=cond, body=body,
                maximum_iterations=length,
                loop_vars=[
                    context_output['presents'],
                    context[:, -1],
                    context,
                ],
                shape_invariants=[
                    tf.TensorShape(model.past_shape(
                        hparams=hparams, batch_size=batch_size)),
                    tf.TensorShape([batch_size]),
                    tf.TensorShape([batch_size, None]),
                ],
                back_prop=False,
            )

        tokens = tokens[:, prefix_len:]

        return tokens
//  [python]
def sample_sequence(*, hparams, length, start_token=None, batch_size=None, context=None, temperature=1, top_k=0):
    if start_token is None:
        assert context is not None, 'Specify exactly one of start_token and context!'
    else:
        assert context is None, 'Specify exactly one of start_token and context!'
        context = tf.fill([batch_size, 1], start_token)

    # get length of prefix
    prefix_len = tf.shape(context)[-1]

    def step(hparams, tokens, past=None):
        lm_output = model.model(
            hparams=hparams, X=tokens, past=past, reuse=tf.AUTO_REUSE)

        logits = lm_output['logits'][:, :, :hparams.n_vocab]
        presents = lm_output['present']
        presents.set_shape(model.past_shape(
            hparams=hparams, batch_size=batch_size))
        return {
            'logits': logits,
            'presents': presents,
        }

    with tf.name_scope('sample_sequence'):
        # Don't feed the last context token -- leave that to the loop below
        # TODO: Would be slightly faster if we called step on the entire context,
        # rather than leaving the last token transformer calculation to the while loop.
        context_output = step(hparams, context[:, :-1])

        def body(past, prev, output):
            next_outputs = step(hparams, prev[:, tf.newaxis], past=past)
            logits = next_outputs['logits'][:, -1, :] / \
                tf.to_float(temperature)
            logits = top_k_logits(logits, k=top_k)
            samples = tf.random.categorical(
                logits, num_samples=1, dtype=tf.int32)
            return [
                tf.concat([past, next_outputs['presents']], axis=-2),
                tf.squeeze(samples, axis=[1]),
                tf.concat([output, samples], axis=1),
            ]

        def cond(*args):
            return True
        with tf.control_dependencies([prefix_len]):
            _, _, tokens = tf.while_loop(
                cond=cond, body=body,
                maximum_iterations=length,
                loop_vars=[
                    context_output['presents'],
                    context[:, -1],
                    context,
                ],
                shape_invariants=[
                    tf.TensorShape(model.past_shape(
                        hparams=hparams, batch_size=batch_size)),
                    tf.TensorShape([batch_size]),
                    tf.TensorShape([batch_size, None]),
                ],
                back_prop=False,
            )

        tokens = tokens[:, prefix_len:]

        return tokens
//  [python]
def categorical_sample(logits, dtype=tf.int32,
                       sample_shape=(), seed=None):
  """Samples from categorical distribution."""
  logits = tf.convert_to_tensor(logits, name="logits")
  event_size = tf.shape(logits)[-1]
  batch_shape_tensor = tf.shape(logits)[:-1]
  def _sample_n(n):
    """Sample vector of categoricals."""
    if logits.shape.ndims == 2:
      logits_2d = logits
    else:
      logits_2d = tf.reshape(logits, [-1, event_size])
    sample_dtype = tf.int64 if logits.dtype.size > 4 else tf.int32
    draws = tf.multinomial(
        logits_2d, n, seed=seed, output_dtype=sample_dtype)
    draws = tf.reshape(
        tf.transpose(draws),
        tf.concat([[n], batch_shape_tensor], 0))
    return tf.cast(draws, dtype)
  return _call_sampler(_sample_n, sample_shape)
//  [python]
def categorical_sample(logits, dtype=tf.int32,
                       sample_shape=(), seed=None):
  """Samples from categorical distribution."""
  logits = tf.convert_to_tensor(logits, name="logits")
  event_size = tf.shape(logits)[-1]
  batch_shape_tensor = tf.shape(logits)[:-1]
  def _sample_n(n):
    """Sample vector of categoricals."""
    if logits.shape.ndims == 2:
      logits_2d = logits
    else:
      logits_2d = tf.reshape(logits, [-1, event_size])
    sample_dtype = tf.int64 if logits.dtype.size > 4 else tf.int32
    draws = tf.random.categorical(
        logits_2d, n, seed=seed, dtype=sample_dtype)
    draws = tf.reshape(
        tf.transpose(draws),
        tf.concat([[n], batch_shape_tensor], 0))
    return tf.cast(draws, dtype)
  return _call_sampler(_sample_n, sample_shape)
//  [python]
def _build_networks(self):
        # Define input placeholders
        self.s = tf.placeholder(tf.float32, shape=[None] + self.state_dim, name='state')
        self.a = tf.placeholder(tf.int32, shape=(None,), name='action')
        self.s_next = tf.placeholder(tf.float32, shape=[None] + self.state_dim, name='next_state')
        self.r = tf.placeholder(tf.float32, shape=(None,), name='reward')
        self.done = tf.placeholder(tf.float32, shape=(None,), name='done_flag')

        # Actor: action probabilities
        self.actor = dense_nn(self.s, self.layer_sizes + [self.act_size], name='actor')
        self.sampled_actions = tf.squeeze(tf.multinomial(self.actor, 1))
        self.actor_proba = tf.nn.softmax(self.actor)
        self.actor_vars = self.scope_vars('actor')

        # Critic: action value (V value)
        self.critic = dense_nn(self.s, self.layer_sizes + [1], name='critic')
        self.critic_next = dense_nn(self.s_next, self.layer_sizes + [1], name='critic', reuse=True)
        self.critic_vars = self.scope_vars('critic')

        # TD target
        self.td_target = self.r + self.gamma * tf.squeeze(self.critic_next) * (1.0 - self.done)
        self.td_error = self.td_target - tf.squeeze(self.critic)
//  [python]
def _build_networks(self):
        # Define input placeholders
        self.s = tf.placeholder(tf.float32, shape=[None] + self.state_dim, name='state')
        self.a = tf.placeholder(tf.int32, shape=(None,), name='action')
        self.s_next = tf.placeholder(tf.float32, shape=[None] + self.state_dim, name='next_state')
        self.r = tf.placeholder(tf.float32, shape=(None,), name='reward')
        self.done = tf.placeholder(tf.float32, shape=(None,), name='done_flag')

        # Actor: action probabilities
        self.actor = dense_nn(self.s, self.layer_sizes + [self.act_size], name='actor')
        self.sampled_actions = tf.squeeze(tf.random.categorical(self.actor, 1))
        self.actor_proba = tf.nn.softmax(self.actor)
        self.actor_vars = self.scope_vars('actor')

        # Critic: action value (V value)
        self.critic = dense_nn(self.s, self.layer_sizes + [1], name='critic')
        self.critic_next = dense_nn(self.s_next, self.layer_sizes + [1], name='critic', reuse=True)
        self.critic_vars = self.scope_vars('critic')

        # TD target
        self.td_target = self.r + self.gamma * tf.squeeze(self.critic_next) * (1.0 - self.done)
        self.td_error = self.td_target - tf.squeeze(self.critic)
//  [python]
def sample_action(self, policy_parameters):
        """ Constructs a symbolic operation for stochastically sampling from the policy
            distribution

            arguments:
                policy_parameters
                    if discrete: logits of a categorical distribution over actions 
                        sy_logits_na: (batch_size, self.ac_dim)
                    if continuous: (mean, log_std) of a Gaussian distribution over actions
                        sy_mean: (batch_size, self.ac_dim)
                        sy_logstd: (self.ac_dim,)

            returns:
                sy_sampled_ac: 
                    if discrete: (batch_size,)
                    if continuous: (batch_size, self.ac_dim)

            Hint: for the continuous case, use the reparameterization trick:
                 The output from a Gaussian distribution with mean 'mu' and std 'sigma' is
        
                      mu + sigma * z,         z ~ N(0, I)
        
                 This reduces the problem to just sampling z. (Hint: use tf.random_normal!)
        """
        # raise NotImplementedError
        if self.discrete:
            sy_logits_na = policy_parameters
            # YOUR_CODE_HERE
            sy_sampled_ac = tf.multinomial(sy_logits_na, num_samples=tf.shape(sy_logits_na)[0])
            sy_sampled_ac = tf.reshape(sy_sampled_ac, [-1])
        else:
            sy_mean, sy_logstd = policy_parameters
            # YOUR_CODE_HERE
            # sy_sampled_ac = sy_mean + sy_logstd*tf.random_normal(tf.shape(sy_logits_na)[0], mean=0.0, stddev=1.0)
            EPSILON=1e-8
            dist = tf.distributions.Normal(loc=sy_mean,scale=tf.exp(sy_logstd+EPSILON),validate_args=True)
            sy_sampled_ac = dist.sample()
        return sy_sampled_ac
//  [python]
def sample_action(self, policy_parameters):
        """ Constructs a symbolic operation for stochastically sampling from the policy
            distribution

            arguments:
                policy_parameters
                    if discrete: logits of a categorical distribution over actions 
                        sy_logits_na: (batch_size, self.ac_dim)
                    if continuous: (mean, log_std) of a Gaussian distribution over actions
                        sy_mean: (batch_size, self.ac_dim)
                        sy_logstd: (self.ac_dim,)

            returns:
                sy_sampled_ac: 
                    if discrete: (batch_size,)
                    if continuous: (batch_size, self.ac_dim)

            Hint: for the continuous case, use the reparameterization trick:
                 The output from a Gaussian distribution with mean 'mu' and std 'sigma' is
        
                      mu + sigma * z,         z ~ N(0, I)
        
                 This reduces the problem to just sampling z. (Hint: use tf.random_normal!)
        """
        # raise NotImplementedError
        if self.discrete:
            sy_logits_na = policy_parameters
            # YOUR_CODE_HERE
            sy_sampled_ac = tf.random.categorical(sy_logits_na, num_samples=tf.shape(sy_logits_na)[0])
            sy_sampled_ac = tf.reshape(sy_sampled_ac, [-1])
        else:
            sy_mean, sy_logstd = policy_parameters
            # YOUR_CODE_HERE
            # sy_sampled_ac = sy_mean + sy_logstd*tf.random_normal(tf.shape(sy_logits_na)[0], mean=0.0, stddev=1.0)
            EPSILON=1e-8
            dist = tf.distributions.Normal(loc=sy_mean,scale=tf.exp(sy_logstd+EPSILON),validate_args=True)
            sy_sampled_ac = dist.sample()
        return sy_sampled_ac