jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def optimize(model, sampler, train, valid):
    """
    Optimize the model. TODO: implement early-stopping
    :param model: model to optimize
    :param sampler: mini-batch sampler
    :param train: train user-item matrix
    :param valid: validation user-item matrix
    :return: None
    """
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())
    if model.feature_projection is not None:
        # initialize item embedding with feature projection
        sess.run(tf.assign(model.item_embeddings, model.feature_projection))

    # sample some users to calculate recall validation
    valid_users = numpy.random.choice(list(set(valid.nonzero()[0])), size=1000, replace=False)

    while True:
        # create evaluator on validation set
        validation_recall = RecallEvaluator(model, train, valid)
        # compute recall on validate set
        valid_recalls = []

        # compute recall in chunks to utilize speedup provided by Tensorflow
        for user_chunk in toolz.partition_all(100, valid_users):
            valid_recalls.extend([validation_recall.eval(sess, user_chunk)])
        print("\nRecall on (sampled) validation set: {}".format(numpy.mean(valid_recalls)))
        # TODO: early stopping based on validation recall

        # train model
        losses = []
        # run n mini-batches
        for _ in tqdm(range(EVALUATION_EVERY_N_BATCHES), desc="Optimizing..."):
            user_pos, neg = sampler.next_batch()
            _, loss = sess.run((model.optimize, model.loss),
                               {model.user_positive_items_pairs: user_pos,
                                model.negative_samples: neg})

            losses.append(loss)

        print("\nTraining loss {}".format(numpy.mean(losses)))
//  [python]
def train_vctk():
    """Summary

    Returns
    -------
    TYPE
        Description
    """
    batch_size = 24
    filter_length = 2
    n_stages = 7
    n_layers_per_stage = 9
    n_hidden = 48
    n_skip = 384
    dataset = vctk.get_dataset()
    it_i = 0
    n_epochs = 1000
    sequence_length = get_sequence_length(n_stages, n_layers_per_stage)
    ckpt_path = 'vctk-wavenet/wavenet_filterlen{}_batchsize{}_sequencelen{}_stages{}_layers{}_hidden{}_skips{}'.format(
        filter_length, batch_size, sequence_length, n_stages,
        n_layers_per_stage, n_hidden, n_skip)
    with tf.Graph().as_default(), tf.Session() as sess:
        net = create_wavenet(
            batch_size=batch_size,
            filter_length=filter_length,
            n_hidden=n_hidden,
            n_skip=n_skip,
            n_stages=n_stages,
            n_layers_per_stage=n_layers_per_stage)
        saver = tf.train.Saver()
        init_op = tf.group(tf.initialize_all_variables(),
                           tf.local_variables_initializer())
        sess.run(init_op)
        if tf.train.latest_checkpoint(ckpt_path) is not None:
            saver.restore(sess, tf.train.latest_checkpoint(ckpt_path))
        batch = vctk.batch_generator
        with tf.variable_scope('optimizer'):
            opt = tf.train.AdamOptimizer(
                learning_rate=0.0002).minimize(net['loss'])
        var_list = [
            v for v in tf.global_variables() if v.name.startswith('optimizer')
        ]
        sess.run(tf.variables_initializer(var_list))
        writer = tf.summary.FileWriter(ckpt_path)
        for epoch_i in range(n_epochs):
            for batch_xs in batch(dataset, batch_size, sequence_length):
                loss, quantized, _ = sess.run(
                    [net['loss'], net['quantized'], opt],
                    feed_dict={net['X']: batch_xs})
                print(loss)
                if it_i % 100 == 0:
                    summary = sess.run(
                        net['summaries'], feed_dict={net['X']: batch_xs})
                    writer.add_summary(summary, it_i)
                    # save
                    saver.save(
                        sess,
                        os.path.join(ckpt_path, 'model.ckpt'),
                        global_step=it_i)
                it_i += 1

    return loss
//  [python]
def train_vctk():
    """Summary

    Returns
    -------
    TYPE
        Description
    """
    batch_size = 24
    filter_length = 2
    n_stages = 7
    n_layers_per_stage = 9
    n_hidden = 48
    n_skip = 384
    dataset = vctk.get_dataset()
    it_i = 0
    n_epochs = 1000
    sequence_length = get_sequence_length(n_stages, n_layers_per_stage)
    ckpt_path = 'vctk-wavenet/wavenet_filterlen{}_batchsize{}_sequencelen{}_stages{}_layers{}_hidden{}_skips{}'.format(
        filter_length, batch_size, sequence_length, n_stages,
        n_layers_per_stage, n_hidden, n_skip)
    with tf.Graph().as_default(), tf.Session() as sess:
        net = create_wavenet(
            batch_size=batch_size,
            filter_length=filter_length,
            n_hidden=n_hidden,
            n_skip=n_skip,
            n_stages=n_stages,
            n_layers_per_stage=n_layers_per_stage)
        saver = tf.train.Saver()
        init_op = tf.group(tf.global_variables_initializer(),
                           tf.local_variables_initializer())
        sess.run(init_op)
        if tf.train.latest_checkpoint(ckpt_path) is not None:
            saver.restore(sess, tf.train.latest_checkpoint(ckpt_path))
        batch = vctk.batch_generator
        with tf.variable_scope('optimizer'):
            opt = tf.train.AdamOptimizer(
                learning_rate=0.0002).minimize(net['loss'])
        var_list = [
            v for v in tf.global_variables() if v.name.startswith('optimizer')
        ]
        sess.run(tf.variables_initializer(var_list))
        writer = tf.summary.FileWriter(ckpt_path)
        for epoch_i in range(n_epochs):
            for batch_xs in batch(dataset, batch_size, sequence_length):
                loss, quantized, _ = sess.run(
                    [net['loss'], net['quantized'], opt],
                    feed_dict={net['X']: batch_xs})
                print(loss)
                if it_i % 100 == 0:
                    summary = sess.run(
                        net['summaries'], feed_dict={net['X']: batch_xs})
                    writer.add_summary(summary, it_i)
                    # save
                    saver.save(
                        sess,
                        os.path.join(ckpt_path, 'model.ckpt'),
                        global_step=it_i)
                it_i += 1

    return loss
//  [python]
def test_librispeech():
    """Summary
    """
    batch_size = 24
    filter_length = 2
    n_stages = 7
    n_layers_per_stage = 9
    n_hidden = 48
    n_skip = 384
    total_length = 16000
    sequence_length = get_sequence_length(n_stages, n_layers_per_stage)
    prime_length = sequence_length
    ckpt_path = 'wavenet/wavenet_filterlen{}_batchsize{}_sequencelen{}_stages{}_layers{}_hidden{}_skips{}/'.format(
        filter_length, batch_size, sequence_length, n_stages,
        n_layers_per_stage, n_hidden, n_skip)

    dataset = librispeech.get_dataset()
    batch = next(
        librispeech.batch_generator(dataset, batch_size, prime_length))[0]

    sess = tf.Session()
    net = create_wavenet(
        batch_size=batch_size,
        filter_length=filter_length,
        n_hidden=n_hidden,
        n_skip=n_skip,
        n_layers_per_stage=n_layers_per_stage,
        n_stages=n_stages,
        shift=False)
    init_op = tf.group(tf.global_variables_initializer(),
                       tf.local_variables_initializer())
    sess.run(init_op)
    saver = tf.train.Saver()
    if tf.train.latest_checkpoint(ckpt_path) is not None:
        saver.restore(sess, tf.train.latest_checkpoint(ckpt_path))
    else:
        print('Could not find checkpoint')

    synth = np.zeros([batch_size, total_length], dtype=np.float32)
    synth[:, :prime_length] = batch

    print('Synthesize...')
    for sample_i in range(0, total_length - prime_length):
        print('{}/{}/{}'.format(sample_i, prime_length, total_length), end='\r')
        probs = sess.run(
            net["probs"],
            feed_dict={net["X"]: synth[:, sample_i:sample_i + sequence_length]})
        idxs = sample_categorical(probs)
        idxs = idxs.reshape((batch_size, sequence_length))
        if sample_i == 0:
            audio = wnu.inv_mu_law_numpy(idxs - 128)
            synth[:, :prime_length] = audio
        else:
            audio = wnu.inv_mu_law_numpy(idxs[:, -1] - 128)
            synth[:, prime_length + sample_i] = audio

    for i in range(batch_size):
        wavfile.write('synthesis-{}.wav'.format(i), 16000, synth[i])
//  [python]
def train_model(hps, logdir):
    """Training."""
    with tf.Graph().as_default():
        with tf.device(tf.train.replica_device_setter(0)):
            with tf.variable_scope("model"):
                model = RealNVP(hps)

            saver = tf.train.Saver(tf.global_variables())

            # Build the summary operation from the last tower summaries.
            summary_op = tf.summary.merge_all()

            # Build an initialization operation to run below.
            init = tf.initialize_all_variables()

            # Start running operations on the Graph. allow_soft_placement must be set to
            # True to build towers on GPU, as some of the ops do not have GPU
            # implementations.
            sess = tf.Session(config=tf.ConfigProto(
                allow_soft_placement=True,
                log_device_placement=True))
            sess.run(init)

            ckpt_state = tf.train.get_checkpoint_state(logdir)
            if ckpt_state and ckpt_state.model_checkpoint_path:
                print("Loading file %s" % ckpt_state.model_checkpoint_path)
                saver.restore(sess, ckpt_state.model_checkpoint_path)

            # Start the queue runners.
            tf.train.start_queue_runners(sess=sess)

            summary_writer = tf.summary.FileWriter(
                logdir,
                graph=sess.graph)

            local_step = 0
            while True:
                fetches = [model.step, model.bit_per_dim, model.train_step]
                # The chief worker evaluates the summaries every 10 steps.
                should_eval_summaries = local_step % 100 == 0
                if should_eval_summaries:
                    fetches += [summary_op]


                start_time = time.time()
                outputs = sess.run(fetches)
                global_step_val = outputs[0]
                loss = outputs[1]
                duration = time.time() - start_time
                assert not numpy.isnan(
                    loss), 'Model diverged with loss = NaN'

                if local_step % 10 == 0:
                    examples_per_sec = hps.batch_size / float(duration)
                    format_str = ('%s: step %d, loss = %.2f '
                                  '(%.1f examples/sec; %.3f '
                                  'sec/batch)')
                    print(format_str % (datetime.now(), global_step_val, loss,
                                        examples_per_sec, duration))

                if should_eval_summaries:
                    summary_str = outputs[-1]
                    summary_writer.add_summary(summary_str, global_step_val)

                # Save the model checkpoint periodically.
                if local_step % 1000 == 0 or (local_step + 1) == FLAGS.train_steps:
                    checkpoint_path = os.path.join(logdir, 'model.ckpt')
                    saver.save(
                        sess,
                        checkpoint_path,
                        global_step=global_step_val)

                if outputs[0] >= FLAGS.train_steps:
                    break

                local_step += 1
//  [python]
def train_model(hps, logdir):
    """Training."""
    with tf.Graph().as_default():
        with tf.device(tf.train.replica_device_setter(0)):
            with tf.variable_scope("model"):
                model = RealNVP(hps)

            saver = tf.train.Saver(tf.global_variables())

            # Build the summary operation from the last tower summaries.
            summary_op = tf.summary.merge_all()

            # Build an initialization operation to run below.
            init = tf.global_variables_initializer()

            # Start running operations on the Graph. allow_soft_placement must be set to
            # True to build towers on GPU, as some of the ops do not have GPU
            # implementations.
            sess = tf.Session(config=tf.ConfigProto(
                allow_soft_placement=True,
                log_device_placement=True))
            sess.run(init)

            ckpt_state = tf.train.get_checkpoint_state(logdir)
            if ckpt_state and ckpt_state.model_checkpoint_path:
                print("Loading file %s" % ckpt_state.model_checkpoint_path)
                saver.restore(sess, ckpt_state.model_checkpoint_path)

            # Start the queue runners.
            tf.train.start_queue_runners(sess=sess)

            summary_writer = tf.summary.FileWriter(
                logdir,
                graph=sess.graph)

            local_step = 0
            while True:
                fetches = [model.step, model.bit_per_dim, model.train_step]
                # The chief worker evaluates the summaries every 10 steps.
                should_eval_summaries = local_step % 100 == 0
                if should_eval_summaries:
                    fetches += [summary_op]


                start_time = time.time()
                outputs = sess.run(fetches)
                global_step_val = outputs[0]
                loss = outputs[1]
                duration = time.time() - start_time
                assert not numpy.isnan(
                    loss), 'Model diverged with loss = NaN'

                if local_step % 10 == 0:
                    examples_per_sec = hps.batch_size / float(duration)
                    format_str = ('%s: step %d, loss = %.2f '
                                  '(%.1f examples/sec; %.3f '
                                  'sec/batch)')
                    print(format_str % (datetime.now(), global_step_val, loss,
                                        examples_per_sec, duration))

                if should_eval_summaries:
                    summary_str = outputs[-1]
                    summary_writer.add_summary(summary_str, global_step_val)

                # Save the model checkpoint periodically.
                if local_step % 1000 == 0 or (local_step + 1) == FLAGS.train_steps:
                    checkpoint_path = os.path.join(logdir, 'model.ckpt')
                    saver.save(
                        sess,
                        checkpoint_path,
                        global_step=global_step_val)

                if outputs[0] >= FLAGS.train_steps:
                    break

                local_step += 1
//  [python]
def train_keywords(self):
        batch = self.iterator.get_next()
        acc, _ = self.predict_keywords(batch)
        with tf.Session(config=self.gpu_config) as sess:
            sess.run(tf.global_variables_initializer())
            sess.run(tf.local_variables_initializer())
            sess.run(tf.tables_initializer())
            self.iterator.switch_to_train_data(sess)

            batchid = 0
            while True:
                try:
                    batchid += 1
                    feed = {tx.global_mode(): tf.estimator.ModeKeys.TRAIN}
                    source_keywords, target_keywords = sess.run([batch['context_text_ids'],
                                                                   batch['keywords_text_ids']], feed_dict=feed)
                    for i in range(len(source_keywords)):
                        for skw_id in source_keywords[i]:
                            if skw_id == 0:
                                break
                            for tkw_id in target_keywords[i]:
                                if skw_id >= 3 and tkw_id >= 3:
                                    tkw = self.config._vocab[tkw_id-4]
                                    if tkw in self.data_config._keywords_candi:
                                        tkw_id = self.data_config._keywords_dict[tkw]
                                        self.pmi_matrix[skw_id][tkw_id] += 1

                except tf.errors.OutOfRangeError:
                    break
            self.pmi_matrix += 0.5
            self.pmi_matrix = self.pmi_matrix / (np.sum(self.pmi_matrix, axis=0) + 1)
            with open(self.config._matrix_save_path,'wb') as f:
                pickle.dump(self.pmi_matrix, f)
//  [python]
def __init__(self, uri):
        """Create a new `TfHubEncoder` object."""
        self._session = tf.Session(graph=tf.Graph())
        with self._session.graph.as_default():
            glog.info("Loading %s model from tensorflow hub", uri)
            embed_fn = tensorflow_hub.Module(uri)
            self._fed_texts = tf.placeholder(shape=[None], dtype=tf.string)
            self._context_embeddings = embed_fn(self._fed_texts)
            init_ops = (
                tf.initialize_all_variables(), tf.tables_initializer())
        glog.info("Initializing graph.")
        self._session.run(init_ops)
//  [python]
def __init__(self, uri):
        """Create a new `TfHubEncoder` object."""
        self._session = tf.Session(graph=tf.Graph())
        with self._session.graph.as_default():
            glog.info("Loading %s model from tensorflow hub", uri)
            embed_fn = tensorflow_hub.Module(uri)
            self._fed_texts = tf.placeholder(shape=[None], dtype=tf.string)
            self._context_embeddings = embed_fn(self._fed_texts)
            init_ops = (
                tf.global_variables_initializer(), tf.tables_initializer())
        glog.info("Initializing graph.")
        self._session.run(init_ops)
//  [python]
def test_keywords(self):
        batch = self.iterator.get_next()
        acc, kws = self.predict_keywords(batch)
        with tf.Session(config=self.gpu_config) as sess:
            sess.run(tf.global_variables_initializer())
            sess.run(tf.local_variables_initializer())
            sess.run(tf.tables_initializer())
            self.iterator.switch_to_test_data(sess)
            cnt_acc, cnt_rec1, cnt_rec3, cnt_rec5 = [], [], [], []
            while True:
                try:
                    feed = {tx.global_mode(): tf.estimator.ModeKeys.PREDICT}
                    acc_, kw_ans, kw_labels = sess.run([acc, kws, batch['keywords_text_ids']], feed_dict=feed)
                    cnt_acc.append(acc_)
                    rec = [0,0,0,0,0]
                    sum_kws = 0
                    for i in range(len(kw_ans)):
                        sum_kws += sum(kw_labels[i] > 3)
                        for j in range(5):
                            if kw_ans[i][j] in kw_labels[i]:
                                for k in range(j, 5):
                                    rec[k] += 1
                    cnt_rec1.append(rec[0]/sum_kws)
                    cnt_rec3.append(rec[2]/sum_kws)
                    cnt_rec5.append(rec[4]/sum_kws)

                except tf.errors.OutOfRangeError:
                    print('test_kw acc@1={:.4f}, rec@1={:.4f}, rec@3={:.4f}, rec@5={:.4f}'.format(
                        np.mean(cnt_acc), np.mean(cnt_rec1), np.mean(cnt_rec3), np.mean(cnt_rec5)))
                    break
//  [python]
def __init__(self, uri):
        """Create a new `ConveRTEncoder` object."""
        self._session = tf.Session(graph=tf.Graph())
        with self._session.graph.as_default():
            glog.info("Loading %s model from tensorflow hub", uri)
            embed_fn = tensorflow_hub.Module(uri)
            self._fed_texts = tf.placeholder(shape=[None], dtype=tf.string)
            self._context_embeddings = embed_fn(
                self._fed_texts, signature="encode_context")
            self._response_embeddings = embed_fn(
                self._fed_texts, signature="encode_response")
            init_ops = (
                tf.initialize_all_variables(), tf.tables_initializer())
        glog.info("Initializing graph.")
        self._session.run(init_ops)
//  [python]
def __init__(self, uri):
        """Create a new `ConveRTEncoder` object."""
        self._session = tf.Session(graph=tf.Graph())
        with self._session.graph.as_default():
            glog.info("Loading %s model from tensorflow hub", uri)
            embed_fn = tensorflow_hub.Module(uri)
            self._fed_texts = tf.placeholder(shape=[None], dtype=tf.string)
            self._context_embeddings = embed_fn(
                self._fed_texts, signature="encode_context")
            self._response_embeddings = embed_fn(
                self._fed_texts, signature="encode_response")
            init_ops = (
                tf.global_variables_initializer(), tf.tables_initializer())
        glog.info("Initializing graph.")
        self._session.run(init_ops)
//  [python]
def testLanguageModelGraph(self):
    train_op, _, _ = graphs.VatxtModel().language_model_training()
    with self.test_session() as sess:
      sess.run(tf.compat.v1.global_variables_initializer())
      tf.train.start_queue_runners(sess)
      sess.run(train_op)
//  [python]
def testLanguageModelGraph(self):
    train_op, _, _ = graphs.VatxtModel().language_model_training()
    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      tf.train.start_queue_runners(sess)
      sess.run(train_op)
//  [python]
def __init__(self, uri):
        """Create a new `USEDualEncoder` object."""
        self._session = tf.Session(graph=tf.Graph())
        with self._session.graph.as_default():
            glog.info("Loading %s model from tensorflow hub", uri)
            embed_fn = tensorflow_hub.Module(uri)
            self._fed_texts = tf.placeholder(shape=[None], dtype=tf.string)
            self._context_embeddings = embed_fn(
                dict(input=self._fed_texts),
                signature="question_encoder",
                as_dict=True,
            )['outputs']
            empty_strings = tf.fill(
                tf.shape(self._fed_texts), ""
            )
            self._response_embeddings = embed_fn(
                dict(input=self._fed_texts, context=empty_strings),
                signature="response_encoder",
                as_dict=True,
            )['outputs']
            init_ops = (
                tf.initialize_all_variables(), tf.tables_initializer())
        glog.info("Initializing graph.")
        self._session.run(init_ops)
//  [python]
def __init__(self, uri):
        """Create a new `USEDualEncoder` object."""
        self._session = tf.Session(graph=tf.Graph())
        with self._session.graph.as_default():
            glog.info("Loading %s model from tensorflow hub", uri)
            embed_fn = tensorflow_hub.Module(uri)
            self._fed_texts = tf.placeholder(shape=[None], dtype=tf.string)
            self._context_embeddings = embed_fn(
                dict(input=self._fed_texts),
                signature="question_encoder",
                as_dict=True,
            )['outputs']
            empty_strings = tf.fill(
                tf.shape(self._fed_texts), ""
            )
            self._response_embeddings = embed_fn(
                dict(input=self._fed_texts, context=empty_strings),
                signature="response_encoder",
                as_dict=True,
            )['outputs']
            init_ops = (
                tf.global_variables_initializer(), tf.tables_initializer())
        glog.info("Initializing graph.")
        self._session.run(init_ops)
//  [python]
def bert_train_fn():
    is_training=True
    hidden_size = 768
    num_labels = 10
    #batch_size=128
    max_seq_length=512
    use_one_hot_embeddings = False
    bert_config = modeling.BertConfig(vocab_size=21128, hidden_size=hidden_size, num_hidden_layers=12,
                                      num_attention_heads=12,intermediate_size=3072)

    input_ids = tf.placeholder(tf.int32, [batch_size, max_seq_length], name="input_ids")
    input_mask = tf.placeholder(tf.int32, [batch_size, max_seq_length], name="input_mask")
    segment_ids = tf.placeholder(tf.int32, [batch_size,max_seq_length],name="segment_ids")
    label_ids = tf.placeholder(tf.float32, [batch_size,num_labels], name="label_ids")
    loss, per_example_loss, logits, probabilities, model = create_model(bert_config, is_training, input_ids, input_mask,
                                                                        segment_ids, label_ids, num_labels,
                                                                        use_one_hot_embeddings)
    # 1. generate or load training/validation/test data. e.g. train:(X,y). X is input_ids,y is labels.

    # 2. train the model by calling create model, get loss
    gpu_config = tf.ConfigProto()
    gpu_config.gpu_options.allow_growth = True
    sess = tf.Session(config=gpu_config)
    sess.run(tf.global_variables_initializer())
    for i in range(1000):
        input_ids_=np.ones((batch_size,max_seq_length),dtype=np.int32)
        input_mask_=np.ones((batch_size,max_seq_length),dtype=np.int32)
        segment_ids_=np.ones((batch_size,max_seq_length),dtype=np.int32)
        label_ids_=np.ones((batch_size,num_labels),dtype=np.float32)
        feed_dict = {input_ids: input_ids_, input_mask: input_mask_,segment_ids:segment_ids_,label_ids:label_ids_}
        loss_ = sess.run([loss], feed_dict)
        print("loss:",loss_)
//  [python]
def testClassifierGraph(self):
    FLAGS.rnn_num_layers = 2
    model = graphs.VatxtModel()
    train_op, _, _ = model.classifier_training()
    # Pretrained vars: embedding + LSTM layers
    self.assertEqual(
        len(model.pretrained_variables), 1 + 2 * FLAGS.rnn_num_layers)
    with self.test_session() as sess:
      sess.run(tf.initialize_all_variables())
      tf.train.start_queue_runners(sess)
      sess.run(train_op)
//  [python]
def testClassifierGraph(self):
    FLAGS.rnn_num_layers = 2
    model = graphs.VatxtModel()
    train_op, _, _ = model.classifier_training()
    # Pretrained vars: embedding + LSTM layers
    self.assertEqual(
        len(model.pretrained_variables), 1 + 2 * FLAGS.rnn_num_layers)
    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      tf.train.start_queue_runners(sess)
      sess.run(train_op)
//  [python]
def initialize(self):
    init_op = tf.initialize_all_variables()
    self.sess.run(init_op)
//  [python]
def initialize(self):
    init_op = tf.global_variables_initializer()
    self.sess.run(init_op)
//  [python]
def save_swa_weight(self):
        average_weights = self.average_weights(self.weight_paths)
        weights = {}
        for k in average_weights.keys():
            weights[k] = tf.get_variable(name=k, initializer=average_weights[k])

        init_op = tf.group(
            tf.initialize_all_variables(),
            tf.local_variables_initializer()
        )

        saver = tf.train.Saver()

        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True

        with tf.Session(config=config) as sess:
            sess.run(init_op)
            saver.save(sess, "../output/trained_weights/{}/swa_{}.ckpt".format(self.cfgs.VERSION, len(self.weight_paths)))
//  [python]
def save_swa_weight(self):
        # img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3])  # is RGB. not BGR
        # img_batch = tf.cast(img_plac, tf.float32)
        #
        # if self.cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']:
        #     img_batch = (img_batch / 255 - tf.constant(self.cfgs.PIXEL_MEAN_)) / tf.constant(self.cfgs.PIXEL_STD)
        # else:
        #     img_batch = img_batch - tf.constant(self.cfgs.PIXEL_MEAN)
        #
        # img_batch = tf.expand_dims(img_batch, axis=0)
        #
        # _ = det_net.build_whole_detection_network(input_img_batch=img_batch)

        average_weights = self.average_weights(self.weight_paths)
        weights = {}
        for k in average_weights.keys():
            weights[k] = tf.get_variable(name=k, initializer=average_weights[k])

        init_op = tf.group(
            tf.global_variables_initializer(),
            tf.local_variables_initializer()
        )

        # average_weights = self.average_weights(weight_paths)

        # model_variables = slim.get_model_variables()
        # for i in range(len(model_variables)):
        #     model_variables[i] = model_variables[i].assign(average_weights[model_variables[i].name.split(':')[0]])

        saver = tf.train.Saver()

        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True

        with tf.Session(config=config) as sess:
            sess.run(init_op)
            # for i in range(len(model_variables)):
            #     model_variables[i].op.run()
            saver.save(sess, "../output/trained_weights/{}/swa_{}.ckpt".format(self.cfgs.VERSION, len(self.weight_paths)))
//  [python]
def main():
    try:
        tf.reset_default_graph()
        sess = tf.InteractiveSession()
        coord = tf.train.Coordinator()

        checkpoint_dir = "checkpoint"
        monitor_dir = "monitors"
        save_path = os.path.join(checkpoint_dir, "model.ckpt")

        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
            print("Directory {} was created".format(checkpoint_dir))

        n_threads = 16
        input_shape = [80, 80, 1]
        output_dim = 3  # {1, 2, 3}
        global_network = A3CNetwork(name="global",
                                    input_shape=input_shape,
                                    output_dim=output_dim)

        thread_list = []
        env_list = []

        for id in range(n_threads):
            env = gym.make("PongDeterministic-v4")

            if id == 0:
                env = gym.wrappers.Monitor(env, monitor_dir, force=True)

            single_agent = Agent(env=env,
                                 session=sess,
                                 coord=coord,
                                 name="thread_{}".format(id),
                                 global_network=global_network,
                                 input_shape=input_shape,
                                 output_dim=output_dim)
            thread_list.append(single_agent)
            env_list.append(env)

        if tf.train.get_checkpoint_state(os.path.dirname(save_path)):
            var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, "global")
            saver = tf.train.Saver(var_list=var_list)
            saver.restore(sess, save_path)
            print("Model restored to global")

        else:
            init = tf.initialize_all_variables()
            sess.run(init)
            print("No model is found")

        for t in thread_list:
            t.start()

        print("Ctrl + C to close")
        coord.wait_for_stop()

    except KeyboardInterrupt:
        var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, "global")
        saver = tf.train.Saver(var_list=var_list)
        saver.save(sess, save_path)
        print('Checkpoint Saved to {}'.format(save_path))

        print("Closing threads")
        coord.request_stop()
        coord.join(thread_list)

        print("Closing environments")
        for env in env_list:
            env.close()

        sess.close()
//  [python]
def main():
    try:
        tf.reset_default_graph()
        sess = tf.InteractiveSession()
        coord = tf.train.Coordinator()

        checkpoint_dir = "checkpoint"
        monitor_dir = "monitors"
        save_path = os.path.join(checkpoint_dir, "model.ckpt")

        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
            print("Directory {} was created".format(checkpoint_dir))

        n_threads = 16
        input_shape = [80, 80, 1]
        output_dim = 3  # {1, 2, 3}
        global_network = A3CNetwork(name="global",
                                    input_shape=input_shape,
                                    output_dim=output_dim)

        thread_list = []
        env_list = []

        for id in range(n_threads):
            env = gym.make("PongDeterministic-v4")

            if id == 0:
                env = gym.wrappers.Monitor(env, monitor_dir, force=True)

            single_agent = Agent(env=env,
                                 session=sess,
                                 coord=coord,
                                 name="thread_{}".format(id),
                                 global_network=global_network,
                                 input_shape=input_shape,
                                 output_dim=output_dim)
            thread_list.append(single_agent)
            env_list.append(env)

        if tf.train.get_checkpoint_state(os.path.dirname(save_path)):
            var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, "global")
            saver = tf.train.Saver(var_list=var_list)
            saver.restore(sess, save_path)
            print("Model restored to global")

        else:
            init = tf.global_variables_initializer()
            sess.run(init)
            print("No model is found")

        for t in thread_list:
            t.start()

        print("Ctrl + C to close")
        coord.wait_for_stop()

    except KeyboardInterrupt:
        var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, "global")
        saver = tf.train.Saver(var_list=var_list)
        saver.save(sess, save_path)
        print('Checkpoint Saved to {}'.format(save_path))

        print("Closing threads")
        coord.request_stop()
        coord.join(thread_list)

        print("Closing environments")
        for env in env_list:
            env.close()

        sess.close()
//  [python]
def test_generator_inference(self):
    """Check one inference step."""
    img_batch = tf.zeros([2, 32, 32, 3])
    model_output, _ = cyclegan.cyclegan_generator_resnet(img_batch)
    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      sess.run(model_output)
//  [python]
def restore_model(self, sess, ckpt_dir, enable_ema=True, export_ckpt=None):
    """Restore variables from checkpoint dir."""
    sess.run(tf.initialize_all_variables())
    checkpoint = tf.train.latest_checkpoint(ckpt_dir)
    if enable_ema:
      ema = tf.train.ExponentialMovingAverage(decay=0.0)
      ema_vars = get_ema_vars()
      var_dict = ema.variables_to_restore(ema_vars)
      ema_assign_op = ema.apply(ema_vars)
    else:
      var_dict = get_ema_vars()
      ema_assign_op = None

    tf.train.get_or_create_global_step()
    sess.run(tf.initialize_all_variables())
    saver = tf.train.Saver(var_dict, max_to_keep=1)
    saver.restore(sess, checkpoint)

    if export_ckpt:
      if ema_assign_op is not None:
        sess.run(ema_assign_op)
      saver = tf.train.Saver(max_to_keep=1, save_relative_paths=True)
      saver.save(sess, export_ckpt)
//  [python]
def restore_model(self, sess, ckpt_dir, enable_ema=True, export_ckpt=None):
    """Restore variables from checkpoint dir."""
    sess.run(tf.global_variables_initializer())
    checkpoint = tf.train.latest_checkpoint(ckpt_dir)
    if enable_ema:
      ema = tf.train.ExponentialMovingAverage(decay=0.0)
      ema_vars = get_ema_vars()
      var_dict = ema.variables_to_restore(ema_vars)
      ema_assign_op = ema.apply(ema_vars)
    else:
      var_dict = get_ema_vars()
      ema_assign_op = None

    tf.train.get_or_create_global_step()
    sess.run(tf.global_variables_initializer())
    saver = tf.train.Saver(var_dict, max_to_keep=1)
    saver.restore(sess, checkpoint)

    if export_ckpt:
      if ema_assign_op is not None:
        sess.run(ema_assign_op)
      saver = tf.train.Saver(max_to_keep=1, save_relative_paths=True)
      saver.save(sess, export_ckpt)
//  [python]
def setup_meta_ops(self):
		cfg = dict({
			'allow_soft_placement': False,
			'log_device_placement': False
		})

		utility = min(self.FLAGS.gpu, 1.)
		if utility > 0.0:
			self.say('GPU mode with {} usage'.format(utility))
			cfg['gpu_options'] = tf.GPUOptions(
				per_process_gpu_memory_fraction = utility)
			cfg['allow_soft_placement'] = True
		else: 
			self.say('Running entirely on CPU')
			cfg['device_count'] = {'GPU': 0}

		if self.FLAGS.train: self.build_train_op()
		
		if self.FLAGS.summary is not None:
			self.summary_op = tf.summary.merge_all()
			self.writer = tf.summary.FileWriter(self.FLAGS.summary + 'train')
		
		self.sess = tf.Session(config = tf.ConfigProto(**cfg))
		self.sess.run(tf.initialize_all_variables())

		if not self.ntrain: return
		self.saver = tf.train.Saver(tf.global_variables(), 
			max_to_keep = self.FLAGS.keep)
		if self.FLAGS.load != 0: self.load_from_ckpt()
		
		if self.FLAGS.summary is not None:
			self.writer.add_graph(self.sess.graph)
//  [python]
def setup_meta_ops(self):
		cfg = dict({
			'allow_soft_placement': False,
			'log_device_placement': False
		})

		utility = min(self.FLAGS.gpu, 1.)
		if utility > 0.0:
			self.say('GPU mode with {} usage'.format(utility))
			cfg['gpu_options'] = tf.GPUOptions(
				per_process_gpu_memory_fraction = utility)
			cfg['allow_soft_placement'] = True
		else: 
			self.say('Running entirely on CPU')
			cfg['device_count'] = {'GPU': 0}

		if self.FLAGS.train: self.build_train_op()
		
		if self.FLAGS.summary is not None:
			self.summary_op = tf.summary.merge_all()
			self.writer = tf.summary.FileWriter(self.FLAGS.summary + 'train')
		
		self.sess = tf.Session(config = tf.ConfigProto(**cfg))
		self.sess.run(tf.global_variables_initializer())

		if not self.ntrain: return
		self.saver = tf.train.Saver(tf.global_variables(), 
			max_to_keep = self.FLAGS.keep)
		if self.FLAGS.load != 0: self.load_from_ckpt()
		
		if self.FLAGS.summary is not None:
			self.writer.add_graph(self.sess.graph)
//  [python]
def train_procedure(self, session, is_continue_train=False):
        """
        训练模型
        """
        if is_continue_train is False:  # 第二条件写完全表明不存在False/True组合
            session.run(tf.global_variables_initializer())

        train_handle = self.get_train_handle(session)
        run_list = [self.train_op, self.loss, self.lr, self.global_step, self.train_score, self.train_label]

        while True:
            try:
                _, loss, lr, step, scores, labels = session.run(run_list, feed_dict={self.handle: train_handle})
                if step % self.args.print_per_steps == 0:
                    acc, acc2, acc5 = self.calc_metrics(scores, labels, self.train_multi_mode)
                    print("step: %s, ls: %.3f, lr: %.5f, acc: %.3f, acc2: %.3f, acc5: %.3f" % (
                        step, loss, lr, acc, acc2, acc5))
                if step % self.args.eval_per_steps == 0:
                    self.test_procedure(session, step)
            except tf.errors.OutOfRangeError as _:
                print("run out of range, epoch done!")
                break
        self.test_procedure(session, step)
//  [python]
def __init__(self, device, model_name, num_actions):
        self.device = device
        self.model_name = model_name
        self.num_actions = num_actions

        self.img_width = Config.IMAGE_WIDTH
        self.img_height = Config.IMAGE_HEIGHT
        self.img_channels = Config.STACKED_FRAMES

        self.learning_rate = Config.LEARNING_RATE_START
        self.beta = Config.BETA_START
        self.log_epsilon = Config.LOG_EPSILON

        self.graph = tf.Graph()
        with self.graph.as_default() as g:
            with tf.device(self.device):
                self._create_graph()

                self.sess = tf.Session(
                    graph=self.graph,
                    config=tf.ConfigProto(
                        allow_soft_placement=True,
                        log_device_placement=False,
                        gpu_options=tf.GPUOptions(allow_growth=True)))
                self.sess.run(tf.initialize_all_variables())

                if Config.TENSORBOARD: self._create_tensor_board()
                if Config.LOAD_CHECKPOINT or Config.SAVE_MODELS:
                    vars = tf.global_variables()
                    self.saver = tf.train.Saver({var.name: var for var in vars}, max_to_keep=0)
//  [python]
def __init__(self, device, model_name, num_actions):
        self.device = device
        self.model_name = model_name
        self.num_actions = num_actions

        self.img_width = Config.IMAGE_WIDTH
        self.img_height = Config.IMAGE_HEIGHT
        self.img_channels = Config.STACKED_FRAMES

        self.learning_rate = Config.LEARNING_RATE_START
        self.beta = Config.BETA_START
        self.log_epsilon = Config.LOG_EPSILON

        self.graph = tf.Graph()
        with self.graph.as_default() as g:
            with tf.device(self.device):
                self._create_graph()

                self.sess = tf.Session(
                    graph=self.graph,
                    config=tf.ConfigProto(
                        allow_soft_placement=True,
                        log_device_placement=False,
                        gpu_options=tf.GPUOptions(allow_growth=True)))
                self.sess.run(tf.global_variables_initializer())

                if Config.TENSORBOARD: self._create_tensor_board()
                if Config.LOAD_CHECKPOINT or Config.SAVE_MODELS:
                    vars = tf.global_variables()
                    self.saver = tf.train.Saver({var.name: var for var in vars}, max_to_keep=0)
//  [python]
def testEvaluation(self):
    batch_size = 2
    num_frames = 64
    height, width = 224, 224
    num_classes = 1000

    eval_inputs = tf.random_uniform((batch_size, num_frames, height, width, 3))
    logits, _ = i3d.i3d(eval_inputs, num_classes,
                        is_training=False)
    predictions = tf.argmax(logits, 1)

    with self.test_session() as sess:
      sess.run(tf.initialize_all_variables())
      output = sess.run(predictions)
      self.assertEquals(output.shape, (batch_size,))
//  [python]
def testEvaluation(self):
    batch_size = 2
    num_frames = 64
    height, width = 224, 224
    num_classes = 1000

    eval_inputs = tf.random_uniform((batch_size, num_frames, height, width, 3))
    logits, _ = i3d.i3d(eval_inputs, num_classes,
                        is_training=False)
    predictions = tf.argmax(logits, 1)

    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      output = sess.run(predictions)
      self.assertEquals(output.shape, (batch_size,))
//  [python]
def train_keywords(self):
        batch = self.iterator.get_next()
        loss, acc, _ = self.predict_keywords(batch)
        op_step = tf.Variable(0, name='op_step')
        train_op = tx.core.get_train_op(loss, global_step=op_step, hparams=self.config.kernel_opt_hparams)
        max_val_acc, stopping_flag = 0, 0
        self.saver = tf.train.Saver()
        with tf.Session(config=self.gpu_config) as sess:
            sess.run(tf.global_variables_initializer())
            sess.run(tf.local_variables_initializer())
            sess.run(tf.tables_initializer())
            for epoch_id in range(self.config._max_epoch):
                self.iterator.switch_to_train_data(sess)
                cur_step = 0
                cnt_acc = []
                while True:
                    try:
                        cur_step += 1
                        feed = {tx.global_mode(): tf.estimator.ModeKeys.TRAIN}
                        loss_, acc_ = sess.run([train_op, acc], feed_dict=feed)
                        cnt_acc.append(acc_)
                        if cur_step % 100 == 0:
                            print('batch {}, loss={}, acc1={}'.format(cur_step, loss_, np.mean(cnt_acc[-100:])))
                    except tf.errors.OutOfRangeError:
                        break

                self.iterator.switch_to_val_data(sess)
                cnt_acc = []
                while True:
                    try:
                        feed = {tx.global_mode(): tf.estimator.ModeKeys.EVAL}
                        acc_ = sess.run(acc, feed_dict=feed)
                        cnt_acc.append(acc_)
                    except tf.errors.OutOfRangeError:
                        mean_acc = np.mean(cnt_acc)
                        if mean_acc > max_val_acc:
                            max_val_acc = mean_acc
                            self.saver.save(sess, self.config._kernel_save_path)
                        else:
                            stopping_flag += 1
                        print('epoch_id {}, valid acc1={}'.format(epoch_id+1, mean_acc))
                        break
                if stopping_flag >= self.config._early_stopping:
                    break
//  [python]
def load_model(self):

        # As discussed in this thread, to use this model in an multi-thread
        # environment (such as a web server). We need to explicitly use the
        # default session:
        # https://github.com/keras-team/keras/issues/2397#issuecomment-254919212
        with self.sess.graph.as_default():
            self.sess.run(tf.initialize_all_variables())
            self.sess.run(tf.local_variables_initializer())
            self.sess.run(tf.tables_initializer())
            # self.sess.run(self.data_iterator)

            print("loading model ", Config.restore_from, Config.dir_model)

            if Config.restore_from:
                self.restore_from_path(Config.restore_from)
            else:
                self.restore_from(Config.dir_model)

            self.summary_writer = tf.summary.FileWriter(
                Config.dir_summary, self.sess.graph, flush_secs=30
            )

            epoch = 0
            while epoch < Config.config_train.max_epochs:
                name = "joint"
                train_op = self.train_ops[name]
                summary_op = self.summary_ops[name]

                step = tf.train.global_step(self.sess, self.global_step)

                self.train_epoch(
                    self.sess,
                    self.summary_writer,
                    "train",
                    train_op,
                    summary_op,
                )

                epoch += 1

                step = tf.train.global_step(self.sess, self.global_step)
                self.save_to(Config.ckpt_model, step)
//  [python]
def load_model(self):

        # As discussed in this thread, to use this model in an multi-thread
        # environment (such as a web server). We need to explicitly use the
        # default session:
        # https://github.com/keras-team/keras/issues/2397#issuecomment-254919212
        with self.sess.graph.as_default():
            self.sess.run(tf.global_variables_initializer())
            self.sess.run(tf.local_variables_initializer())
            self.sess.run(tf.tables_initializer())
            # self.sess.run(self.data_iterator)

            print("loading model ", Config.restore_from, Config.dir_model)

            if Config.restore_from:
                self.restore_from_path(Config.restore_from)
            else:
                self.restore_from(Config.dir_model)

            self.summary_writer = tf.summary.FileWriter(
                Config.dir_summary, self.sess.graph, flush_secs=30
            )

            epoch = 0
            while epoch < Config.config_train.max_epochs:
                name = "joint"
                train_op = self.train_ops[name]
                summary_op = self.summary_ops[name]

                step = tf.train.global_step(self.sess, self.global_step)

                self.train_epoch(
                    self.sess,
                    self.summary_writer,
                    "train",
                    train_op,
                    summary_op,
                )

                epoch += 1

                step = tf.train.global_step(self.sess, self.global_step)
                self.save_to(Config.ckpt_model, step)