jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def train(self, train_config):
        # Define training graphs
        self._build_training_graph(train_config)

        self.lr = tf.train.exponential_decay(train_config["learning_rate"],
                                             self.global_step,
                                             train_config["decay_steps"],
                                             train_config["lr_decay"],
                                             staircase=True)
        self.optim_op = tf.train.AdamOptimizer(self.lr,
                                               beta1=train_config["beta1"]) \
            .minimize(self.loss,
                      var_list=tf.trainable_variables(),
                      global_step=self.global_step)

        ginit_op = tf.global_variables_initializer()
        linit_op = tf.local_variables_initializer()
        self.sess.run([ginit_op, linit_op])

        self.writer = tf.summary.FileWriter(train_config["logs_dir"],
                                            self.sess.graph)
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=self.sess, coord=coord)

        start_time = time.time()
        for v in tf.trainable_variables():
            print(v)
        if train_config["continue_train"] and \
                self.load(train_config["checkpoint_dir"]):
            print('[*] Load Success')
        elif train_config["restore_dir"] and \
                self.load(train_config["restore_dir"],
                          tf.train.Saver(var_list=tf.trainable_variables())):
            self.sess.run(self.global_step.assign(0))
            print('[*] Restore success')
        else:
            print('Training from scratch.')
        try:
            while not coord.should_stop():
                _, loss_sum_str = self.sess.run([self.optim_op, self.loss_sum])
                global_step = self.sess.run(self.global_step)
                self.writer.add_summary(loss_sum_str, global_step)

                if global_step % 100 == 0:
                    # Write image summary.
                    img_sum_str, img_comp_str, img_orig_str = \
                            self.sess.run([self.image_sum,
                                           self.image_comp_sum,
                                           self.image_orig_comp_sum])
                    self.writer.add_summary(img_sum_str, global_step)
                    self.writer.add_summary(img_comp_str, global_step)
                    self.writer.add_summary(img_orig_str, global_step)

                elapsed_time = time.time() - start_time
                print ("Steps: %2d time: %4.4f (%4.4f steps/sec)" % (
                    global_step, elapsed_time,
                    float(global_step) / elapsed_time))

                if np.mod(global_step, train_config["save_freq"]) == 2:
                    self.save(train_config["checkpoint_dir"], global_step)

        except tf.errors.OutOfRangeError:
            print('Done Training.')
        finally:
            coord.request_stop()
            coord.join(threads)
//  [python]
def step_train(train_op,loss,accuracy,
               train_images_batch,train_labels_batch,train_nums,train_log_step,
               val_images_batch,val_labels_batch,val_nums,val_log_step,
               snapshot_prefix,snapshot):
    '''
    循环迭代训练过程
    :param train_op: 训练op
    :param loss:     loss函数
    :param accuracy: 准确率函数
    :param train_images_batch: 训练images数据
    :param train_labels_batch: 训练labels数据
    :param train_nums:         总训练数据
    :param train_log_step:   训练log显示间隔
    :param val_images_batch: 验证images数据
    :param val_labels_batch: 验证labels数据
    :param val_nums:         总验证数据
    :param val_log_step:     验证log显示间隔
    :param snapshot_prefix: 模型保存的路径
    :param snapshot:        模型保存间隔
    :return: None
    '''
    saver = tf.train.Saver()
    max_acc = 0.0
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        sess.run(tf.local_variables_initializer())
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)
        for i in range(max_steps + 1):
            batch_input_images, batch_input_labels = sess.run([train_images_batch, train_labels_batch])
            _, train_loss = sess.run([train_op, loss], feed_dict={input_images: batch_input_images,
                                                                  input_labels: batch_input_labels,
                                                                  keep_prob: 0.8, is_training: True})
            # train测试(这里仅测试训练集的一个batch)
            if i % train_log_step == 0:
                train_acc = sess.run(accuracy, feed_dict={input_images: batch_input_images,
                                                          input_labels: batch_input_labels,
                                                          keep_prob: 1.0, is_training: False})
                print("%s: Step [%d]  train Loss : %f, training accuracy :  %g" % (
                datetime.now(), i, train_loss, train_acc))

            # val测试(测试全部val数据)
            if i % val_log_step == 0:
                mean_loss, mean_acc = net_evaluation(sess, loss, accuracy, val_images_batch, val_labels_batch, val_nums)
                print("%s: Step [%d]  val Loss : %f, val accuracy :  %g" % (datetime.now(), i, mean_loss, mean_acc))

            # 模型保存:每迭代snapshot次或者最后一次保存模型
            if (i % snapshot == 0 and i > 0) or i == max_steps:
                print('-----save:{}-{}'.format(snapshot_prefix, i))
                saver.save(sess, snapshot_prefix, global_step=i)
            # 保存val准确率最高的模型
            if mean_acc > max_acc and mean_acc > 0.7:
                max_acc = mean_acc
                path = os.path.dirname(snapshot_prefix)
                best_models = os.path.join(path, 'best_models_{}_{:.4f}.ckpt'.format(i, max_acc))
                print('------save:{}'.format(best_models))
                saver.save(sess, best_models)

        coord.request_stop()
        coord.join(threads)
//  [python]
def __init__(
            self,
            n_actions,
            n_features,
            learning_rate=0.005,
            reward_decay=0.9,
            e_greedy=0.9,
            replace_target_iter=500,
            memory_size=10000,
            batch_size=32,
            e_greedy_increment=None,
            hidden=[100, 50],
            output_graph=False,
            sess=None,
    ):
        self.n_actions = n_actions
        self.n_features = n_features
        self.lr = learning_rate
        self.gamma = reward_decay
        self.epsilon_max = e_greedy
        self.replace_target_iter = replace_target_iter
        self.memory_size = memory_size
        self.batch_size = batch_size
        self.hidden = hidden
        self.epsilon_increment = e_greedy_increment
        self.epsilon = 0.5 if e_greedy_increment is not None else self.epsilon_max

        self.learn_step_counter = 0
        self._build_net()
        self.memory = Memory(capacity=memory_size)

        if sess is None:
            self.sess = tf.Session()
            self.sess.run(tf.initialize_all_variables())
        else:
            self.sess = sess

        if output_graph:
            tf.summary.FileWriter("logs/", self.sess.graph)

        self.cost_his = []
//  [python]
def __init__(
            self,
            n_actions,
            n_features,
            learning_rate=0.005,
            reward_decay=0.9,
            e_greedy=0.9,
            replace_target_iter=500,
            memory_size=10000,
            batch_size=32,
            e_greedy_increment=None,
            hidden=[100, 50],
            output_graph=False,
            sess=None,
    ):
        self.n_actions = n_actions
        self.n_features = n_features
        self.lr = learning_rate
        self.gamma = reward_decay
        self.epsilon_max = e_greedy
        self.replace_target_iter = replace_target_iter
        self.memory_size = memory_size
        self.batch_size = batch_size
        self.hidden = hidden
        self.epsilon_increment = e_greedy_increment
        self.epsilon = 0.5 if e_greedy_increment is not None else self.epsilon_max

        self.learn_step_counter = 0
        self._build_net()
        self.memory = Memory(capacity=memory_size)

        if sess is None:
            self.sess = tf.Session()
            self.sess.run(tf.global_variables_initializer())
        else:
            self.sess = sess

        if output_graph:
            tf.summary.FileWriter("logs/", self.sess.graph)

        self.cost_his = []
//  [python]
def inference(sess, dataset, passes=1,
              dataset_mask=False, dataset_selection=False,
              min_opacity=.15, max_opacity=.4):
    # Data sources
    next_image, iterator_init = dataset()
    image_shape = next_image.shape.as_list()
    if dataset_mask:
        next_mask, iterator_init = dataset_mask()
        next_mask = next_mask[:, :, :, 0:1]
    else:
        next_mask = batch_masks(
            None, image_shape[1], image_shape[2], min_opacity, max_opacity)
    if dataset_selection:
        next_selection, iterator_init = dataset_selection()

    # Model
    images_p = tf.placeholder(tf.float32, shape=[None] + image_shape[1:])
    mask_p = tf.placeholder(tf.float32, shape=[None] + image_shape[1:3] + [1])
    selection_p = tf.placeholder(tf.float32, shape=[None] + image_shape[1:3] + [1])

    image_w = tf.clip_by_value(images_p - mask_p, 0, 1)
    selection_conv = selection_margin(mask_p, 4)
    predictions = model(image_w, False) * selection_p
    gen_mask = tf.clip_by_value(tf.abs(predictions), 0, 1)
    reconstruction = tf.clip_by_value(image_w + predictions, 0, 1)
    accuracy = get_accuracy(reconstruction, images_p)

    # Inference
    sess.run(tf.initialize_all_variables())
    sess.run(tf.tables_initializer())
    iterator_init(sess)

    try:
        saver = tf.train.Saver()
        saver.restore(sess, "/tmp/model.ckpt")
    except Exception as e:
        print('An error occurred when trying to restore : %s' % e)

    # Pass 1
    batch = sess.run([next_image, next_mask])
    if dataset_selection:
        selection = sess.run(tf.expand_dims(next_selection[:, :, :, 0], 3))
    else:
        selection = sess.run(selection_conv, feed_dict={mask_p: batch[1]})

    feed_dict = {images_p: batch[0][:, :, :, 0:3],
                 mask_p: batch[1][:batch[0].shape[0]],
                 selection_p: selection[:batch[0].shape[0]]}
    images = sess.run([image_w, reconstruction, gen_mask, accuracy], feed_dict=feed_dict)
    results = [images[0], images[1]]
    print('Mean accuracy %.3f%%' % (images[3] * 100))

    # Pass 2
    for _ in range(1, passes):
        reconstruction1 = images[1]
        feed_dict = {images_p: reconstruction1,
                     mask_p: np.zeros(list(reconstruction1.shape[:-1]) + [1])}
        images = sess.run([image_w, reconstruction, gen_mask], feed_dict=feed_dict)
        results += [images[1]]

    # Reformat
    results += [images[2]]
    images = [unstandardize(x) for x in results]
    return images
//  [python]
def inference(sess, dataset, passes=1,
              dataset_mask=False, dataset_selection=False,
              min_opacity=.15, max_opacity=.4):
    # Data sources
    next_image, iterator_init = dataset()
    image_shape = next_image.shape.as_list()
    if dataset_mask:
        next_mask, iterator_init = dataset_mask()
        next_mask = next_mask[:, :, :, 0:1]
    else:
        next_mask = batch_masks(
            None, image_shape[1], image_shape[2], min_opacity, max_opacity)
    if dataset_selection:
        next_selection, iterator_init = dataset_selection()

    # Model
    images_p = tf.placeholder(tf.float32, shape=[None] + image_shape[1:])
    mask_p = tf.placeholder(tf.float32, shape=[None] + image_shape[1:3] + [1])
    selection_p = tf.placeholder(tf.float32, shape=[None] + image_shape[1:3] + [1])

    image_w = tf.clip_by_value(images_p - mask_p, 0, 1)
    selection_conv = selection_margin(mask_p, 4)
    predictions = model(image_w, False) * selection_p
    gen_mask = tf.clip_by_value(tf.abs(predictions), 0, 1)
    reconstruction = tf.clip_by_value(image_w + predictions, 0, 1)
    accuracy = get_accuracy(reconstruction, images_p)

    # Inference
    sess.run(tf.global_variables_initializer())
    sess.run(tf.tables_initializer())
    iterator_init(sess)

    try:
        saver = tf.train.Saver()
        saver.restore(sess, "/tmp/model.ckpt")
    except Exception as e:
        print('An error occurred when trying to restore : %s' % e)

    # Pass 1
    batch = sess.run([next_image, next_mask])
    if dataset_selection:
        selection = sess.run(tf.expand_dims(next_selection[:, :, :, 0], 3))
    else:
        selection = sess.run(selection_conv, feed_dict={mask_p: batch[1]})

    feed_dict = {images_p: batch[0][:, :, :, 0:3],
                 mask_p: batch[1][:batch[0].shape[0]],
                 selection_p: selection[:batch[0].shape[0]]}
    images = sess.run([image_w, reconstruction, gen_mask, accuracy], feed_dict=feed_dict)
    results = [images[0], images[1]]
    print('Mean accuracy %.3f%%' % (images[3] * 100))

    # Pass 2
    for _ in range(1, passes):
        reconstruction1 = images[1]
        feed_dict = {images_p: reconstruction1,
                     mask_p: np.zeros(list(reconstruction1.shape[:-1]) + [1])}
        images = sess.run([image_w, reconstruction, gen_mask], feed_dict=feed_dict)
        results += [images[1]]

    # Reformat
    results += [images[2]]
    images = [unstandardize(x) for x in results]
    return images
//  [python]
def run(self):
		n_users, n_items, n_train, n_test = self.data_counter()
		train_items, test_set = self.data_generator()
		plain_adj, norm_adj, mean_adj = self.get_adj_mat()
		model = self.model(
			lr=self.args['lr'],
			embed_size=self.args['embed_size'],
			batch_size=self.args['batch_size'],
			layer_size=self.args['layer_size'],
			regs=self.args['regs'],
			n_users=n_users,
			n_items=n_items,
			adj=plain_adj)

		config = tf.ConfigProto()
		sess = tf.Session(config=config)
		sess.run(tf.compat.v1.global_variables_initializer())
		best_res = 0.


		for epoch in range(self.args['epoch']):
			t1 = time()
			loss, mf_loss, emb_loss, reg_loss = 0., 0., 0., 0.
			n_batch = n_train // self.args['batch_size'] + 1
			for idx in range(n_batch):
				users, pos_items, neg_items = self.sample(self.args['batch_size'])
				_, batch_loss, batch_mf_loss, batch_emb_loss, batch_reg_loss = sess.run(
					[model.opt, model.loss, model.mf_loss, model.emb_loss, model.reg_loss],
					feed_dict={model.users: users, 
							model.pos_items: pos_items,
							model.node_dropout: eval(self.args['node_dropout']), 
							model.mess_dropout: eval(self.args['mess_dropout']),
							model.neg_items: neg_items})
				loss += batch_loss
				mf_loss += batch_mf_loss
				emb_loss += batch_emb_loss
				reg_loss += batch_reg_loss

			if np.isnan(loss) == True:
				print('ERROR: loss is nan.')
				sys.exit()

			if (epoch + 1) % 10 != 0:
				perf_str = 'Epoch %d [%.1fs]: train==[%.5f=%.5f + %.5f]' % (epoch, time() - t1, loss, mf_loss, reg_loss)
				print(perf_str)
				continue

			t2 = time()
			users_to_test = list(test_set.keys())
			res = self.evaluate(sess, model, users_to_test, self.args['batch_size'], drop_flag=True)

			t3 = time()

			perf_str = 'Epoch %d [%.1fs + %.1fs]: train==[%.5f=%.5f + %.5f + %.5f], recall=[%.5f, %.5f], ' \
					'precision=[%.5f, %.5f], hit=[%.5f, %.5f], ndcg=[%.5f, %.5f]' % \
					(epoch, t2 - t1, t3 - t2, loss, mf_loss, emb_loss, reg_loss, res['recall'][0], res['recall'][-1],
					res['precision'][0], res['precision'][-1], res['hit_ratio'][0], res['hit_ratio'][-1],
					res['ndcg'][0], res['ndcg'][-1])
			print(perf_str)

			best_res, stopping_step, should_stop = self.early_stopping(
				res['recall'][0], best_res, stopping_step, expected_order='acc', flag_step=5)

			if should_stop == True:
				break

			if res['recall'][0] == best_res and args.save_flag == 1:
				save_saver.save(sess, self.args['model_dir'] + '/weights', global_step=epoch)
				print('save the weights in path: ', self.args['model_dir'])
//  [python]
def run(self):
		n_users, n_items, n_train, n_test = self.data_counter()
		train_items, test_set = self.data_generator()
		plain_adj, norm_adj, mean_adj = self.get_adj_mat()
		model = self.model(
			lr=self.args['lr'],
			embed_size=self.args['embed_size'],
			batch_size=self.args['batch_size'],
			layer_size=self.args['layer_size'],
			regs=self.args['regs'],
			n_users=n_users,
			n_items=n_items,
			adj=plain_adj)

		config = tf.ConfigProto()
		sess = tf.Session(config=config)
		sess.run(tf.global_variables_initializer())
		best_res = 0.


		for epoch in range(self.args['epoch']):
			t1 = time()
			loss, mf_loss, emb_loss, reg_loss = 0., 0., 0., 0.
			n_batch = n_train // self.args['batch_size'] + 1
			for idx in range(n_batch):
				users, pos_items, neg_items = self.sample(self.args['batch_size'])
				_, batch_loss, batch_mf_loss, batch_emb_loss, batch_reg_loss = sess.run(
					[model.opt, model.loss, model.mf_loss, model.emb_loss, model.reg_loss],
					feed_dict={model.users: users, 
							model.pos_items: pos_items,
							model.node_dropout: eval(self.args['node_dropout']), 
							model.mess_dropout: eval(self.args['mess_dropout']),
							model.neg_items: neg_items})
				loss += batch_loss
				mf_loss += batch_mf_loss
				emb_loss += batch_emb_loss
				reg_loss += batch_reg_loss

			if np.isnan(loss) == True:
				print('ERROR: loss is nan.')
				sys.exit()

			if (epoch + 1) % 10 != 0:
				perf_str = 'Epoch %d [%.1fs]: train==[%.5f=%.5f + %.5f]' % (epoch, time() - t1, loss, mf_loss, reg_loss)
				print(perf_str)
				continue

			t2 = time()
			users_to_test = list(test_set.keys())
			res = self.evaluate(sess, model, users_to_test, self.args['batch_size'], drop_flag=True)

			t3 = time()

			perf_str = 'Epoch %d [%.1fs + %.1fs]: train==[%.5f=%.5f + %.5f + %.5f], recall=[%.5f, %.5f], ' \
					'precision=[%.5f, %.5f], hit=[%.5f, %.5f], ndcg=[%.5f, %.5f]' % \
					(epoch, t2 - t1, t3 - t2, loss, mf_loss, emb_loss, reg_loss, res['recall'][0], res['recall'][-1],
					res['precision'][0], res['precision'][-1], res['hit_ratio'][0], res['hit_ratio'][-1],
					res['ndcg'][0], res['ndcg'][-1])
			print(perf_str)

			best_res, stopping_step, should_stop = self.early_stopping(
				res['recall'][0], best_res, stopping_step, expected_order='acc', flag_step=5)

			if should_stop == True:
				break

			if res['recall'][0] == best_res and args.save_flag == 1:
				save_saver.save(sess, self.args['model_dir'] + '/weights', global_step=epoch)
				print('save the weights in path: ', self.args['model_dir'])
//  [python]
def train(sess, dataset, min_opacity=.15, max_opacity=.4):
    global_step = tf.Variable(0, name='global_step', trainable=False)
    training = tf.placeholder(tf.bool, shape=[])
    with tf.device('/cpu:0'):
        next_image, iterator_inits = dataset_split(dataset, .8)

    masks = batch_masks(
        global_step, next_image.shape.as_list()[1], next_image.shape.as_list()[2],
        min_opacity, max_opacity)
    image_w = tf.clip_by_value(next_image - masks, 0, 1)
    predictions = model(image_w, training) * selection_margin(masks, 4)
    tf.summary.image('masks', predictions)

    # Define loss
    image_mask = -(image_w - next_image)  # Mask after application on the image
    abs_loss = tf.losses.absolute_difference(
        predictions, image_mask, loss_collection=None)**.5
    tf.losses.add_loss(abs_loss)
    loss = tf.losses.get_total_loss(True)
    tf.summary.scalar('loss', loss)

    # Optimizer
    learning_rate = tf.train.polynomial_decay(
        FLAGS.learning_rate, global_step,
        decay_steps=60000, end_learning_rate=.0005)
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        train_op = tf.train.AdamOptimizer(learning_rate).minimize(
            loss,
            global_step=global_step)

    # Training loop
    sess.run(tf.initialize_all_variables())
    sess.run(iterator_inits[0])

    saver = tf.train.Saver()
    summaries = tf.summary.merge_all()
    train_writer = tf.summary.FileWriter(os.path.join(FLAGS.logdir, 'train'), sess.graph)
    val_writer = tf.summary.FileWriter(os.path.join(FLAGS.logdir, 'val'), sess.graph)

    for i in range(1, 2 if DEBUG else int(1e6)):
        if DEBUG:
            start_time = time()
            _, loss_, predictions_ = sess.run([train_op, loss, predictions],
                                              feed_dict={training: True})
            batch_time = 1000 * (time() - start_time) / FLAGS.batch_size
            print('Time %dms, Loss %f' % (batch_time, loss_))
            continue

        _, summaries_, global_step_ = sess.run(
            [train_op, summaries, global_step], feed_dict={training: True})
        train_writer.add_summary(summaries_, global_step_)

        # Save model
        if i % 2000 == 0:
            path = saver.save(sess, "/tmp/model.ckpt")
            print(i, 'Saving at', path)
            sess.run(iterator_inits[1])  # switch to validation dataset
            while True:
                try:
                    _, summaries_ = sess.run([loss, summaries],
                                             feed_dict={training: False})
                    val_writer.add_summary(summaries_, global_step_)
                except tf.errors.OutOfRangeError:
                    break
            sess.run(iterator_inits[0])
    return
//  [python]
def train(sess, dataset, min_opacity=.15, max_opacity=.4):
    global_step = tf.Variable(0, name='global_step', trainable=False)
    training = tf.placeholder(tf.bool, shape=[])
    with tf.device('/cpu:0'):
        next_image, iterator_inits = dataset_split(dataset, .8)

    masks = batch_masks(
        global_step, next_image.shape.as_list()[1], next_image.shape.as_list()[2],
        min_opacity, max_opacity)
    image_w = tf.clip_by_value(next_image - masks, 0, 1)
    predictions = model(image_w, training) * selection_margin(masks, 4)
    tf.summary.image('masks', predictions)

    # Define loss
    image_mask = -(image_w - next_image)  # Mask after application on the image
    abs_loss = tf.losses.absolute_difference(
        predictions, image_mask, loss_collection=None)**.5
    tf.losses.add_loss(abs_loss)
    loss = tf.losses.get_total_loss(True)
    tf.summary.scalar('loss', loss)

    # Optimizer
    learning_rate = tf.train.polynomial_decay(
        FLAGS.learning_rate, global_step,
        decay_steps=60000, end_learning_rate=.0005)
    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
        train_op = tf.train.AdamOptimizer(learning_rate).minimize(
            loss,
            global_step=global_step)

    # Training loop
    sess.run(tf.global_variables_initializer())
    sess.run(iterator_inits[0])

    saver = tf.train.Saver()
    summaries = tf.summary.merge_all()
    train_writer = tf.summary.FileWriter(os.path.join(FLAGS.logdir, 'train'), sess.graph)
    val_writer = tf.summary.FileWriter(os.path.join(FLAGS.logdir, 'val'), sess.graph)

    for i in range(1, 2 if DEBUG else int(1e6)):
        if DEBUG:
            start_time = time()
            _, loss_, predictions_ = sess.run([train_op, loss, predictions],
                                              feed_dict={training: True})
            batch_time = 1000 * (time() - start_time) / FLAGS.batch_size
            print('Time %dms, Loss %f' % (batch_time, loss_))
            continue

        _, summaries_, global_step_ = sess.run(
            [train_op, summaries, global_step], feed_dict={training: True})
        train_writer.add_summary(summaries_, global_step_)

        # Save model
        if i % 2000 == 0:
            path = saver.save(sess, "/tmp/model.ckpt")
            print(i, 'Saving at', path)
            sess.run(iterator_inits[1])  # switch to validation dataset
            while True:
                try:
                    _, summaries_ = sess.run([loss, summaries],
                                             feed_dict={training: False})
                    val_writer.add_summary(summaries_, global_step_)
                except tf.errors.OutOfRangeError:
                    break
            sess.run(iterator_inits[0])
    return
//  [python]
def run_training():
    """Train MNIST for a number of steps."""

    # Tell TensorFlow that the model will be built into the default Graph.
    with tf.Graph().as_default():
        # Input images and labels.
        image_batch, label_batch = inputs(
            train=True, batch_size=FLAGS.batch_size, num_epochs=FLAGS.num_epochs
        )

        # Build a Graph that computes predictions from the inference model.
        logits = mnist.inference(image_batch, FLAGS.hidden1, FLAGS.hidden2)

        # Add to the Graph the loss calculation.
        loss = mnist.loss(logits, label_batch)

        # Add to the Graph operations that train the model.
        train_op = mnist.training(loss, FLAGS.learning_rate)

        # The op for initializing the variables.
        init_op = tf.group(tf.initialize_all_variables(), tf.local_variables_initializer())

        # Create a session for running operations in the Graph.
        with tf.Session() as sess:
            # Initialize the variables (the trained variables and the
            # epoch counter).
            sess.run(init_op)
            try:
                step = 0
                while True:  # Train until OutOfRangeError
                    start_time = time.time()

                    # Run one step of the model.  The return values are
                    # the activations from the `train_op` (which is
                    # discarded) and the `loss` op.  To inspect the values
                    # of your ops or variables, you may include them in
                    # the list passed to sess.run() and the value tensors
                    # will be returned in the tuple from the call.
                    _, loss_value = sess.run([train_op, loss])

                    duration = time.time() - start_time

                    # Print an overview fairly often.
                    if step % 100 == 0:
                        print("Step %d: loss = %.2f (%.3f sec)" % (step, loss_value, duration))
                    step += 1
            except tf.errors.OutOfRangeError:
                print("Done training for %d epochs, %d steps." % (FLAGS.num_epochs, step))
//  [python]
def run_training():
    """Train MNIST for a number of steps."""

    # Tell TensorFlow that the model will be built into the default Graph.
    with tf.Graph().as_default():
        # Input images and labels.
        image_batch, label_batch = inputs(
            train=True, batch_size=FLAGS.batch_size, num_epochs=FLAGS.num_epochs
        )

        # Build a Graph that computes predictions from the inference model.
        logits = mnist.inference(image_batch, FLAGS.hidden1, FLAGS.hidden2)

        # Add to the Graph the loss calculation.
        loss = mnist.loss(logits, label_batch)

        # Add to the Graph operations that train the model.
        train_op = mnist.training(loss, FLAGS.learning_rate)

        # The op for initializing the variables.
        init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())

        # Create a session for running operations in the Graph.
        with tf.Session() as sess:
            # Initialize the variables (the trained variables and the
            # epoch counter).
            sess.run(init_op)
            try:
                step = 0
                while True:  # Train until OutOfRangeError
                    start_time = time.time()

                    # Run one step of the model.  The return values are
                    # the activations from the `train_op` (which is
                    # discarded) and the `loss` op.  To inspect the values
                    # of your ops or variables, you may include them in
                    # the list passed to sess.run() and the value tensors
                    # will be returned in the tuple from the call.
                    _, loss_value = sess.run([train_op, loss])

                    duration = time.time() - start_time

                    # Print an overview fairly often.
                    if step % 100 == 0:
                        print("Step %d: loss = %.2f (%.3f sec)" % (step, loss_value, duration))
                    step += 1
            except tf.errors.OutOfRangeError:
                print("Done training for %d epochs, %d steps." % (FLAGS.num_epochs, step))
//  [python]
def __init__(self, batchsize=64,
                 epochs=20, learning_rate=1e-4, 
                 dropout_rate=0.5,
                 shuffle=True, random_seed=None):
        np.random.seed(random_seed)
        self.batchsize = batchsize
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.dropout_rate = dropout_rate
        self.shuffle = shuffle
                
        g = tf.Graph()
        with g.as_default():
            ## set random-seed:
            tf.set_random_seed(random_seed)
            
            ## build the network:
            self.build()

            ## initializer
            self.init_op = tf.initialize_all_variables()

            ## saver
            self.saver = tf.train.Saver()
            
        ## create a session
        self.sess = tf.Session(graph=g)
//  [python]
def __init__(self, batchsize=64,
                 epochs=20, learning_rate=1e-4, 
                 dropout_rate=0.5,
                 shuffle=True, random_seed=None):
        np.random.seed(random_seed)
        self.batchsize = batchsize
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.dropout_rate = dropout_rate
        self.shuffle = shuffle
                
        g = tf.Graph()
        with g.as_default():
            ## set random-seed:
            tf.set_random_seed(random_seed)
            
            ## build the network:
            self.build()

            ## initializer
            self.init_op =                 tf.global_variables_initializer()

            ## saver
            self.saver = tf.train.Saver()
            
        ## create a session
        self.sess = tf.Session(graph=g)
//  [python]
def buildModel(self):
        print ('training...')
        iteration = 0
        with tf.Session() as sess:
            init = tf.initialize_all_variables()
            sess.run(init)
            # train the model until converged
            for epoch in range(self.maxIter):
                user_idx,item_idx,neg_item_idx = self.next_batch()
                _,loss = sess.run([self.train,self.total_loss],feed_dict={self.u_idx: user_idx, self.v_idx: item_idx, self.neg_idx:neg_item_idx})
                print ('iteration:', epoch, 'loss:',loss)

                self.P = sess.run(self.U)
                self.Q = sess.run(self.V)
                self.ranking_performance()
            # start adversarial training
            for epoch in range(self.advEpoch):
                user_idx,item_idx,neg_item_idx = self.next_batch()
                sess.run([self.update_U, self.update_V],
                         feed_dict={self.u_idx: user_idx, self.v_idx: item_idx, self.neg_idx: neg_item_idx})
                _,loss = sess.run([self.train_adv,self.loss_adv],feed_dict={self.u_idx: user_idx, self.v_idx: item_idx, self.neg_idx:neg_item_idx})
                print ('iteration:', epoch, 'loss:',loss)
                self.P = sess.run(self.U)
                self.Q = sess.run(self.V)
                self.ranking_performance()
//  [python]
def buildModel(self):
        print ('training...')
        iteration = 0
        with tf.Session() as sess:
            init = tf.global_variables_initializer()
            sess.run(init)
            # train the model until converged
            for epoch in range(self.maxIter):
                user_idx,item_idx,neg_item_idx = self.next_batch()
                _,loss = sess.run([self.train,self.total_loss],feed_dict={self.u_idx: user_idx, self.v_idx: item_idx, self.neg_idx:neg_item_idx})
                print ('iteration:', epoch, 'loss:',loss)

                self.P = sess.run(self.U)
                self.Q = sess.run(self.V)
                self.ranking_performance()
            # start adversarial training
            for epoch in range(self.advEpoch):
                user_idx,item_idx,neg_item_idx = self.next_batch()
                sess.run([self.update_U, self.update_V],
                         feed_dict={self.u_idx: user_idx, self.v_idx: item_idx, self.neg_idx: neg_item_idx})
                _,loss = sess.run([self.train_adv,self.loss_adv],feed_dict={self.u_idx: user_idx, self.v_idx: item_idx, self.neg_idx:neg_item_idx})
                print ('iteration:', epoch, 'loss:',loss)
                self.P = sess.run(self.U)
                self.Q = sess.run(self.V)
                self.ranking_performance()
//  [python]
def train(sess, training_set, validation_set=None,
          initialize=True, epochs=20, shuffle=True,
          dropout=0.5, random_seed=None):

    X_data = np.array(training_set[0])
    y_data = np.array(training_set[1])
    training_loss = []

    ## initialize variables
    if initialize:
        sess.run(tf.initialize_all_variables())

    np.random.seed(random_seed) # for shuflling in batch_generator
    for epoch in range(1, epochs+1):
        batch_gen = batch_generator(
                        X_data, y_data, 
                        shuffle=shuffle)
        avg_loss = 0.0
        for i,(batch_x,batch_y) in enumerate(batch_gen):
            feed = {'tf_x:0': batch_x, 
                    'tf_y:0': batch_y, 
                    'fc_keep_prob:0': dropout}
            loss, _ = sess.run(
                    ['cross_entropy_loss:0', 'train_op'],
                    feed_dict=feed)
            avg_loss += loss

        training_loss.append(avg_loss / (i+1))
        print('Epoch %02d Training Avg. Loss: %7.3f' % (
            epoch, avg_loss), end=' ')
        if validation_set is not None:
            feed = {'tf_x:0': validation_set[0],
                    'tf_y:0': validation_set[1],
                    'fc_keep_prob:0':1.0}
            valid_acc = sess.run('accuracy:0', feed_dict=feed)
            print(' Validation Acc: %7.3f' % valid_acc)
        else:
            print()
//  [python]
def train(sess, training_set, validation_set=None,
          initialize=True, epochs=20, shuffle=True,
          dropout=0.5, random_seed=None):

    X_data = np.array(training_set[0])
    y_data = np.array(training_set[1])
    training_loss = []

    ## initialize variables
    if initialize:
        sess.run(tf.global_variables_initializer())

    np.random.seed(random_seed) # for shuflling in batch_generator
    for epoch in range(1, epochs+1):
        batch_gen = batch_generator(
                        X_data, y_data, 
                        shuffle=shuffle)
        avg_loss = 0.0
        for i,(batch_x,batch_y) in enumerate(batch_gen):
            feed = {'tf_x:0': batch_x, 
                    'tf_y:0': batch_y, 
                    'fc_keep_prob:0': dropout}
            loss, _ = sess.run(
                    ['cross_entropy_loss:0', 'train_op'],
                    feed_dict=feed)
            avg_loss += loss

        training_loss.append(avg_loss / (i+1))
        print('Epoch %02d Training Avg. Loss: %7.3f' % (
            epoch, avg_loss), end=' ')
        if validation_set is not None:
            feed = {'tf_x:0': validation_set[0],
                    'tf_y:0': validation_set[1],
                    'fc_keep_prob:0':1.0}
            valid_acc = sess.run('accuracy:0', feed_dict=feed)
            print(' Validation Acc: %7.3f' % valid_acc)
        else:
            print()
//  [python]
def get_distributed_session_creator(server):
    """
    Args:
       server (tf.train.Server):

    Returns:
        tf.train.SessionCreator
    """

    server_def = server.server_def
    is_chief = (server_def.job_name == 'worker') and (server_def.task_index == 0)

    init_op = tf.initialize_all_variables()
    local_init_op = tf.local_variables_initializer()
    ready_op = tf.report_uninitialized_variables()
    ready_for_local_init_op = tf.report_uninitialized_variables(tf.global_variables())
    sm = tf.train.SessionManager(
        local_init_op=local_init_op,
        ready_op=ready_op,
        ready_for_local_init_op=ready_for_local_init_op,
        graph=tf.get_default_graph())

    # to debug wrong variable collection
    # from pprint import pprint
    # print("GLOBAL:")
    # pprint([(k.name, k.device) for k in tf.global_variables()])
    # print("LOCAL:")
    # pprint([(k.name, k.device) for k in tf.local_variables()])

    class _Creator(tf.train.SessionCreator):
        def create_session(self):
            if is_chief:
                return sm.prepare_session(master=server.target, init_op=init_op)
            else:
                tf.logging.set_verbosity(tf.logging.INFO)   # print message about uninitialized vars
                ret = sm.wait_for_session(master=server.target)
                tf.logging.set_verbosity(tf.logging.WARN)
                return ret

    return _Creator()
//  [python]
def get_distributed_session_creator(server):
    """
    Args:
       server (tf.train.Server):

    Returns:
        tf.train.SessionCreator
    """

    server_def = server.server_def
    is_chief = (server_def.job_name == 'worker') and (server_def.task_index == 0)

    init_op = tf.global_variables_initializer()
    local_init_op = tf.local_variables_initializer()
    ready_op = tf.report_uninitialized_variables()
    ready_for_local_init_op = tf.report_uninitialized_variables(tf.global_variables())
    sm = tf.train.SessionManager(
        local_init_op=local_init_op,
        ready_op=ready_op,
        ready_for_local_init_op=ready_for_local_init_op,
        graph=tf.get_default_graph())

    # to debug wrong variable collection
    # from pprint import pprint
    # print("GLOBAL:")
    # pprint([(k.name, k.device) for k in tf.global_variables()])
    # print("LOCAL:")
    # pprint([(k.name, k.device) for k in tf.local_variables()])

    class _Creator(tf.train.SessionCreator):
        def create_session(self):
            if is_chief:
                return sm.prepare_session(master=server.target, init_op=init_op)
            else:
                tf.logging.set_verbosity(tf.logging.INFO)   # print message about uninitialized vars
                ret = sm.wait_for_session(master=server.target)
                tf.logging.set_verbosity(tf.logging.WARN)
                return ret

    return _Creator()
//  [python]
def run_benchmark(model_details, args, find_graph_def):
    tf_config = create_tf_config(args)
    graph = initialize_graph(model_details, args, find_graph_def)
    run_options = tf_v1.RunOptions(trace_level=tf_v1.RunOptions.FULL_TRACE)
    run_metadata = tf_v1.RunMetadata()

    if args.save_graph:
        # write the real benchmark graph to local
        model_dir = os.path.dirname(os.path.abspath(model_detail['model_dir']))
        out_graph_file = os.path.join(model_dir, 'runtime_graph.pb')
        write_graph(graph.as_graph_def(), out_graph_file)
        print("********** save runtime graph at {}".format(out_graph_file))
 
    with tf_v1.Session(config=tf_config, graph=graph) as sess:
        input_dict = {graph.get_tensor_by_name("g/" + out_name + ":0"): model_details['input'][out_name]
                       for out_name in model_details['input']}
        output_dict = {out_name: graph.get_tensor_by_name("g/" + out_name + ":0")
                       for out_name in model_details['output']}

        sess.run(tf_v1.initialize_all_variables())

        total_time = 0.0
        reps_done = 0
        for rep in range(args.num_iter):
            # sess run
            start = time.time()
            if args.profile:
                _ = sess.run(output_dict, options=run_options, run_metadata=run_metadata)
            else:
                _ = sess.run(output_dict, feed_dict=input_dict)
            end = time.time()
            delta = end - start
            print("Iteration: {}, inference time: {} sec".format(rep, delta), flush=True)

            if rep >= args.num_warmup and rep < (args.num_iter - args.num_warmup):
                total_time += delta
                reps_done += 1

        avg_time = total_time / reps_done
        latency = avg_time * 1000
        throughput = 1.0 / avg_time * args.batch_size
        print('Batch size = %d' % args.batch_size)
        print("Latency: {:.3f} ms".format(latency))
        print("Throughput: {:.2f} fps".format(throughput))

        # Logging to a file
        log_file = open("log.txt", "a")
        log_file.write("Throughput: " + str(throughput) + "\n")
//  [python]
def run_benchmark(model_details, args, find_graph_def):
    tf_config = create_tf_config(args)
    graph = initialize_graph(model_details, args, find_graph_def)
    run_options = tf_v1.RunOptions(trace_level=tf_v1.RunOptions.FULL_TRACE)
    run_metadata = tf_v1.RunMetadata()

    if args.save_graph:
        # write the real benchmark graph to local
        model_dir = os.path.dirname(os.path.abspath(model_detail['model_dir']))
        out_graph_file = os.path.join(model_dir, 'runtime_graph.pb')
        write_graph(graph.as_graph_def(), out_graph_file)
        print("********** save runtime graph at {}".format(out_graph_file))
 
    with tf_v1.Session(config=tf_config, graph=graph) as sess:
        input_dict = {graph.get_tensor_by_name("g/" + out_name + ":0"): model_details['input'][out_name]
                       for out_name in model_details['input']}
        output_dict = {out_name: graph.get_tensor_by_name("g/" + out_name + ":0")
                       for out_name in model_details['output']}

        sess.run(tf_v1.global_variables_initializer())

        total_time = 0.0
        reps_done = 0
        for rep in range(args.num_iter):
            # sess run
            start = time.time()
            if args.profile:
                _ = sess.run(output_dict, options=run_options, run_metadata=run_metadata)
            else:
                _ = sess.run(output_dict, feed_dict=input_dict)
            end = time.time()
            delta = end - start
            print("Iteration: {}, inference time: {} sec".format(rep, delta), flush=True)

            if rep >= args.num_warmup and rep < (args.num_iter - args.num_warmup):
                total_time += delta
                reps_done += 1

        avg_time = total_time / reps_done
        latency = avg_time * 1000
        throughput = 1.0 / avg_time * args.batch_size
        print('Batch size = %d' % args.batch_size)
        print("Latency: {:.3f} ms".format(latency))
        print("Throughput: {:.2f} fps".format(throughput))

        # Logging to a file
        log_file = open("log.txt", "a")
        log_file.write("Throughput: " + str(throughput) + "\n")
//  [python]
def testAtrousFullyConvolutionalValues(self):
    """Verify dense feature extraction with atrous convolution."""
    nominal_stride = 32
    for output_stride in [4, 8, 16, 32, None]:
      with slim.arg_scope(resnet_utils.resnet_arg_scope()):
        with tf.Graph().as_default():
          with self.test_session() as sess:
            tf.set_random_seed(0)
            inputs = create_test_input(2, 81, 81, 3)
            # Dense feature extraction followed by subsampling.
            output, _ = self._resnet_small(inputs, None,
                                           is_training=False,
                                           global_pool=False,
                                           output_stride=output_stride)
            if output_stride is None:
              factor = 1
            else:
              factor = nominal_stride // output_stride
            output = resnet_utils.subsample(output, factor)
            # Make the two networks use the same weights.
            tf.get_variable_scope().reuse_variables()
            # Feature extraction at the nominal network rate.
            expected, _ = self._resnet_small(inputs, None,
                                             is_training=False,
                                             global_pool=False)
            sess.run(tf.initialize_all_variables())
            self.assertAllClose(output.eval(), expected.eval(),
                                atol=1e-4, rtol=1e-4)
//  [python]
def testAtrousFullyConvolutionalValues(self):
    """Verify dense feature extraction with atrous convolution."""
    nominal_stride = 32
    for output_stride in [4, 8, 16, 32, None]:
      with slim.arg_scope(resnet_utils.resnet_arg_scope()):
        with tf.Graph().as_default():
          with self.test_session() as sess:
            tf.set_random_seed(0)
            inputs = create_test_input(2, 81, 81, 3)
            # Dense feature extraction followed by subsampling.
            output, _ = self._resnet_small(inputs, None,
                                           is_training=False,
                                           global_pool=False,
                                           output_stride=output_stride)
            if output_stride is None:
              factor = 1
            else:
              factor = nominal_stride // output_stride
            output = resnet_utils.subsample(output, factor)
            # Make the two networks use the same weights.
            tf.get_variable_scope().reuse_variables()
            # Feature extraction at the nominal network rate.
            expected, _ = self._resnet_small(inputs, None,
                                             is_training=False,
                                             global_pool=False)
            sess.run(tf.global_variables_initializer())
            self.assertAllClose(output.eval(), expected.eval(),
                                atol=1e-4, rtol=1e-4)