jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def eval_(
    args,
    log_dir,
    env_id,
    seed,
    model_file,
    **kwargs) :
    with open(str(Path(log_dir)/'args.txt')) as f:
        from argparse import Namespace
        model_args = eval(f.read())

    np.random.seed(seed)
    tf.random.set_random_seed(seed)

    # Env
    env = wrap_deepmind(make_atari(env_id), # Noop randomization/skip frame of 4 is applied.
                        episode_life=False,
                        clip_rewards=False,
                        frame_stack=True,
                        scale=False)
    num_ac = env.action_space.n

    # Neural Episodic Controller
    nec = NEC(num_ac,model_args.p,model_args.embed_size,1e-5,0.,0.,dnd_params={
        'maxlen': model_args.memory_len,
        'seed': model_args.seed,
        'cores': model_args.cores, # #cores for KD-Tree
        'trees': model_args.trees, # #trees for KD-Tree
    })

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

    nec.restore(os.path.join(log_dir,model_file))

    while(True):
        ep_len, acc_r = _run(env,nec,'test.mp4')
        print(f'ep_len: {ep_len} acc_r: {acc_r}')

        input()
//  [python]
def test_restore_map_for_classification_ckpt(self):
    # Define mock tensorflow classification graph and save variables.
    test_graph_classification = tf.Graph()
    with test_graph_classification.as_default():
      image = tf.placeholder(dtype=tf.float32, shape=[1, 20, 20, 3])
      with tf.variable_scope('mock_model'):
        net = slim.conv2d(image, num_outputs=3, kernel_size=1, scope='layer1')
        slim.conv2d(net, num_outputs=3, kernel_size=1, scope='layer2')

      init_op = tf.initialize_all_variables()
      saver = tf.train.Saver()
      save_path = self.get_temp_dir()
      with self.test_session(graph=test_graph_classification) as sess:
        sess.run(init_op)
        saved_model_path = saver.save(sess, save_path)

    # Create tensorflow detection graph and load variables from
    # classification checkpoint.
    test_graph_detection = tf.Graph()
    with test_graph_detection.as_default():
      model = self._build_model(
          is_training=False, number_of_stages=2, second_stage_batch_size=6)

      inputs_shape = (2, 20, 20, 3)
      inputs = tf.to_float(tf.random_uniform(
          inputs_shape, minval=0, maxval=255, dtype=tf.int32))
      preprocessed_inputs, true_image_shapes = model.preprocess(inputs)
      prediction_dict = model.predict(preprocessed_inputs, true_image_shapes)
      model.postprocess(prediction_dict, true_image_shapes)
      var_map = model.restore_map(fine_tune_checkpoint_type='classification')
      self.assertIsInstance(var_map, dict)
      saver = tf.train.Saver(var_map)
      with self.test_session(graph=test_graph_classification) as sess:
        saver.restore(sess, saved_model_path)
        for var in sess.run(tf.report_uninitialized_variables()):
          self.assertNotIn(model.first_stage_feature_extractor_scope, var)
          self.assertNotIn(model.second_stage_feature_extractor_scope, var)
//  [python]
def test_restore_map_for_classification_ckpt(self):
    # Define mock tensorflow classification graph and save variables.
    test_graph_classification = tf.Graph()
    with test_graph_classification.as_default():
      image = tf.placeholder(dtype=tf.float32, shape=[1, 20, 20, 3])
      with tf.variable_scope('mock_model'):
        net = slim.conv2d(image, num_outputs=3, kernel_size=1, scope='layer1')
        slim.conv2d(net, num_outputs=3, kernel_size=1, scope='layer2')

      init_op = tf.global_variables_initializer()
      saver = tf.train.Saver()
      save_path = self.get_temp_dir()
      with self.test_session(graph=test_graph_classification) as sess:
        sess.run(init_op)
        saved_model_path = saver.save(sess, save_path)

    # Create tensorflow detection graph and load variables from
    # classification checkpoint.
    test_graph_detection = tf.Graph()
    with test_graph_detection.as_default():
      model = self._build_model(
          is_training=False, number_of_stages=2, second_stage_batch_size=6)

      inputs_shape = (2, 20, 20, 3)
      inputs = tf.to_float(tf.random_uniform(
          inputs_shape, minval=0, maxval=255, dtype=tf.int32))
      preprocessed_inputs, true_image_shapes = model.preprocess(inputs)
      prediction_dict = model.predict(preprocessed_inputs, true_image_shapes)
      model.postprocess(prediction_dict, true_image_shapes)
      var_map = model.restore_map(fine_tune_checkpoint_type='classification')
      self.assertIsInstance(var_map, dict)
      saver = tf.train.Saver(var_map)
      with self.test_session(graph=test_graph_classification) as sess:
        saver.restore(sess, saved_model_path)
        for var in sess.run(tf.report_uninitialized_variables()):
          self.assertNotIn(model.first_stage_feature_extractor_scope, var)
          self.assertNotIn(model.second_stage_feature_extractor_scope, var)
//  [python]
def check_extract_features_raises_error_with_invalid_image_size(
      self, image_height, image_width, depth_multiplier, pad_to_multiple,
      use_keras=False):
    preprocessed_inputs = tf.placeholder(tf.float32, (4, None, None, 3))
    feature_maps = self._extract_features(preprocessed_inputs,
                                          depth_multiplier,
                                          pad_to_multiple,
                                          use_keras=use_keras)
    test_preprocessed_image = np.random.rand(4, image_height, image_width, 3)
    with self.test_session() as sess:
      sess.run(tf.initialize_all_variables())
      with self.assertRaises(tf.errors.InvalidArgumentError):
        sess.run(feature_maps,
                 feed_dict={preprocessed_inputs: test_preprocessed_image})
//  [python]
def check_extract_features_raises_error_with_invalid_image_size(
      self, image_height, image_width, depth_multiplier, pad_to_multiple,
      use_keras=False):
    preprocessed_inputs = tf.placeholder(tf.float32, (4, None, None, 3))
    feature_maps = self._extract_features(preprocessed_inputs,
                                          depth_multiplier,
                                          pad_to_multiple,
                                          use_keras=use_keras)
    test_preprocessed_image = np.random.rand(4, image_height, image_width, 3)
    with self.test_session() as sess:
      sess.run(tf.global_variables_initializer())
      with self.assertRaises(tf.errors.InvalidArgumentError):
        sess.run(feature_maps,
                 feed_dict={preprocessed_inputs: test_preprocessed_image})
//  [python]
def test_predict_gives_valid_anchors_in_training_mode_first_stage_only(self):
    test_graph = tf.Graph()
    with test_graph.as_default():
      model = self._build_model(
          is_training=True, number_of_stages=1, second_stage_batch_size=2)
      batch_size = 2
      height = 10
      width = 12
      input_image_shape = (batch_size, height, width, 3)
      _, true_image_shapes = model.preprocess(tf.zeros(input_image_shape))
      preprocessed_inputs = tf.placeholder(
          dtype=tf.float32, shape=(batch_size, None, None, 3))
      prediction_dict = model.predict(preprocessed_inputs, true_image_shapes)

      expected_output_keys = set([
          'rpn_box_predictor_features', 'rpn_features_to_crop', 'image_shape',
          'rpn_box_encodings', 'rpn_objectness_predictions_with_background',
          'anchors'])
      # At training time, anchors that exceed image bounds are pruned.  Thus
      # the `expected_num_anchors` in the above inference mode test is now
      # a strict upper bound on the number of anchors.
      num_anchors_strict_upper_bound = height * width * 3 * 3

      init_op = tf.initialize_all_variables()
      with self.test_session(graph=test_graph) as sess:
        sess.run(init_op)
        prediction_out = sess.run(prediction_dict,
                                  feed_dict={
                                      preprocessed_inputs:
                                      np.zeros(input_image_shape)
                                  })

        self.assertEqual(set(prediction_out.keys()), expected_output_keys)
        self.assertAllEqual(prediction_out['image_shape'], input_image_shape)

        # Check that anchors have less than the upper bound and
        # are clipped to window.
        anchors = prediction_out['anchors']
        self.assertTrue(len(anchors.shape) == 2 and anchors.shape[1] == 4)
        num_anchors_out = anchors.shape[0]
        self.assertLess(num_anchors_out, num_anchors_strict_upper_bound)

        self.assertTrue(np.all(np.greater_equal(anchors, 0)))
        self.assertTrue(np.all(np.less_equal(anchors[:, 0], height)))
        self.assertTrue(np.all(np.less_equal(anchors[:, 1], width)))
        self.assertTrue(np.all(np.less_equal(anchors[:, 2], height)))
        self.assertTrue(np.all(np.less_equal(anchors[:, 3], width)))

        self.assertAllEqual(prediction_out['rpn_box_encodings'].shape,
                            (batch_size, num_anchors_out, 4))
        self.assertAllEqual(
            prediction_out['rpn_objectness_predictions_with_background'].shape,
            (batch_size, num_anchors_out, 2))
//  [python]
def test_predict_gives_valid_anchors_in_training_mode_first_stage_only(self):
    test_graph = tf.Graph()
    with test_graph.as_default():
      model = self._build_model(
          is_training=True, number_of_stages=1, second_stage_batch_size=2)
      batch_size = 2
      height = 10
      width = 12
      input_image_shape = (batch_size, height, width, 3)
      _, true_image_shapes = model.preprocess(tf.zeros(input_image_shape))
      preprocessed_inputs = tf.placeholder(
          dtype=tf.float32, shape=(batch_size, None, None, 3))
      prediction_dict = model.predict(preprocessed_inputs, true_image_shapes)

      expected_output_keys = set([
          'rpn_box_predictor_features', 'rpn_features_to_crop', 'image_shape',
          'rpn_box_encodings', 'rpn_objectness_predictions_with_background',
          'anchors'])
      # At training time, anchors that exceed image bounds are pruned.  Thus
      # the `expected_num_anchors` in the above inference mode test is now
      # a strict upper bound on the number of anchors.
      num_anchors_strict_upper_bound = height * width * 3 * 3

      init_op = tf.global_variables_initializer()
      with self.test_session(graph=test_graph) as sess:
        sess.run(init_op)
        prediction_out = sess.run(prediction_dict,
                                  feed_dict={
                                      preprocessed_inputs:
                                      np.zeros(input_image_shape)
                                  })

        self.assertEqual(set(prediction_out.keys()), expected_output_keys)
        self.assertAllEqual(prediction_out['image_shape'], input_image_shape)

        # Check that anchors have less than the upper bound and
        # are clipped to window.
        anchors = prediction_out['anchors']
        self.assertTrue(len(anchors.shape) == 2 and anchors.shape[1] == 4)
        num_anchors_out = anchors.shape[0]
        self.assertLess(num_anchors_out, num_anchors_strict_upper_bound)

        self.assertTrue(np.all(np.greater_equal(anchors, 0)))
        self.assertTrue(np.all(np.less_equal(anchors[:, 0], height)))
        self.assertTrue(np.all(np.less_equal(anchors[:, 1], width)))
        self.assertTrue(np.all(np.less_equal(anchors[:, 2], height)))
        self.assertTrue(np.all(np.less_equal(anchors[:, 3], width)))

        self.assertAllEqual(prediction_out['rpn_box_encodings'].shape,
                            (batch_size, num_anchors_out, 4))
        self.assertAllEqual(
            prediction_out['rpn_objectness_predictions_with_background'].shape,
            (batch_size, num_anchors_out, 2))
//  [python]
def tf_baseline_basic():
    import tensorflow as tf
    import cntk.contrib.crosstalk.crosstalk_tensorflow as crtf
    ci = crtf.instance

    tf.reset_default_graph()
    
    p1 = tf.get_variable("param1", initializer=param1, dtype=tf.float32)
    p2 = tf.get_variable("param2", initializer=param2, dtype=tf.float32)
    ci.watch(p1, 'p1', var_type=crtf.TrainableType)
    ci.watch({'param1':p1, 'param2':p2}, 'p1_p2', var_type=crtf.DictTrainableType)

    with tf.Session() as sess:
        sess.run(tf.initialize_all_variables())
        ci.set_workdir(workdir)
        ci.set_data(sess, None)
        ci.fetch('p1', save=True)
        ci.fetch('p1_p2', save=True)
        ci.reset()
        sess.close()
//  [python]
def tf_baseline_basic():
    import tensorflow as tf
    import cntk.contrib.crosstalk.crosstalk_tensorflow as crtf
    ci = crtf.instance

    tf.reset_default_graph()
    
    p1 = tf.get_variable("param1", initializer=param1, dtype=tf.float32)
    p2 = tf.get_variable("param2", initializer=param2, dtype=tf.float32)
    ci.watch(p1, 'p1', var_type=crtf.TrainableType)
    ci.watch({'param1':p1, 'param2':p2}, 'p1_p2', var_type=crtf.DictTrainableType)

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        ci.set_workdir(workdir)
        ci.set_data(sess, None)
        ci.fetch('p1', save=True)
        ci.fetch('p1_p2', save=True)
        ci.reset()
        sess.close()
//  [python]
def load_checkpoint_v1(self):
        ckpt_path = os.path.dirname(self._tf_file)
        latest_ckpt = tf.train.latest_checkpoint(ckpt_path)
        saver = tf.train.import_meta_graph(latest_ckpt + ".meta")
        with tf.Session() as session:
            session.run(
                [
                    tf.global_variables_initializer(),
                    tf.local_variables_initializer()
                ]
            )
            saver.restore(session, latest_ckpt)
            graph_def = session.graph.as_graph_def(add_shapes=True)
        frozen_graph = freeze_graph.freeze_graph_with_def_protos(
            input_graph_def=graph_def,
            input_saver_def=None,
            input_checkpoint=latest_ckpt,
            output_node_names=self._outputs,
            restore_op_name="",
            filename_tensor_name="",
            output_graph=None,
            clear_devices=True,
            initializer_nodes=""
        )
        return frozen_graph
//  [python]
def test_output_size_nn_upsample_conv(self):
    batch_size = 2
    height, width = 256, 256
    num_outputs = 4

    images = tf.ones((batch_size, height, width, 3))
    with tf.contrib.framework.arg_scope(pix2pix.pix2pix_arg_scope()):
      logits, _ = pix2pix.pix2pix_generator(
          images, num_outputs, blocks=self._reduced_default_blocks(),
          upsample_method='nn_upsample_conv')

    with self.test_session() as session:
      session.run(tf.initialize_all_variables())
      np_outputs = session.run(logits)
      self.assertListEqual([batch_size, height, width, num_outputs],
                           list(np_outputs.shape))
//  [python]
def test_output_size_nn_upsample_conv(self):
    batch_size = 2
    height, width = 256, 256
    num_outputs = 4

    images = tf.ones((batch_size, height, width, 3))
    with tf.contrib.framework.arg_scope(pix2pix.pix2pix_arg_scope()):
      logits, _ = pix2pix.pix2pix_generator(
          images, num_outputs, blocks=self._reduced_default_blocks(),
          upsample_method='nn_upsample_conv')

    with self.test_session() as session:
      session.run(tf.global_variables_initializer())
      np_outputs = session.run(logits)
      self.assertListEqual([batch_size, height, width, num_outputs],
                           list(np_outputs.shape))
//  [python]
def launch(self):
    with self._graph.as_default():
      self.merged = tf.summary.merge_all()
      init = tf.initialize_all_variables()
    self.session = tf.Session(graph=self._graph)
    self.session.run(init)
//  [python]
def launch(self):
    with self._graph.as_default():
      self.merged = tf.summary.merge_all()
      init = tf.global_variables_initializer()
    self.session = tf.Session(graph=self._graph)
    self.session.run(init)
//  [python]
def test_output_size_conv2d_transpose(self):
    batch_size = 2
    height, width = 256, 256
    num_outputs = 4

    images = tf.ones((batch_size, height, width, 3))
    with tf.contrib.framework.arg_scope(pix2pix.pix2pix_arg_scope()):
      logits, _ = pix2pix.pix2pix_generator(
          images, num_outputs, blocks=self._reduced_default_blocks(),
          upsample_method='conv2d_transpose')

    with self.test_session() as session:
      session.run(tf.initialize_all_variables())
      np_outputs = session.run(logits)
      self.assertListEqual([batch_size, height, width, num_outputs],
                           list(np_outputs.shape))
//  [python]
def test_output_size_conv2d_transpose(self):
    batch_size = 2
    height, width = 256, 256
    num_outputs = 4

    images = tf.ones((batch_size, height, width, 3))
    with tf.contrib.framework.arg_scope(pix2pix.pix2pix_arg_scope()):
      logits, _ = pix2pix.pix2pix_generator(
          images, num_outputs, blocks=self._reduced_default_blocks(),
          upsample_method='conv2d_transpose')

    with self.test_session() as session:
      session.run(tf.global_variables_initializer())
      np_outputs = session.run(logits)
      self.assertListEqual([batch_size, height, width, num_outputs],
                           list(np_outputs.shape))
//  [python]
def _createTestTrainModel(self, m_creator, hparams, sess):
    train_mode = tf.contrib.learn.ModeKeys.TRAIN
    train_iterator, src_vocab_table, tgt_vocab_table = (
        common_test_utils.create_test_iterator(hparams, train_mode))
    train_m = m_creator(
        hparams,
        train_mode,
        train_iterator,
        src_vocab_table,
        tgt_vocab_table,
        scope='dynamic_seq2seq')
    sess.run(tf.global_variables_initializer())
    sess.run(tf.tables_initializer())
    sess.run(train_iterator.initializer)
    return train_m
//  [python]
def real_signal():
    spf = wave.open('helloworld.wav', 'r')

    #Extract Raw Audio from Wav File
    # If you right-click on the file and go to "Get Info", you can see:
    # sampling rate = 16000 Hz
    # bits per sample = 16
    # The first is quantization in time
    # The second is quantization in amplitude
    # We also do this for images!
    # 2^16 = 65536 is how many different sound levels we have
    signal = spf.readframes(-1)
    signal = np.fromstring(signal, 'Int16')
    T = len(signal)
    signal = (signal - signal.mean()) / signal.std()

    # signal needs to be of shape N x T(n) x D
    signals = signal.reshape(1, T, 1)

    hmm = HMM(3, 3, 1)
    hmm.init_random(signals)

    init = tf.initialize_all_variables()
    session = tf.InteractiveSession()
    session.run(init)
    hmm.set_session(session)

    hmm.fit(signals, max_iter=30)
//  [python]
def real_signal():
    spf = wave.open('helloworld.wav', 'r')

    #Extract Raw Audio from Wav File
    # If you right-click on the file and go to "Get Info", you can see:
    # sampling rate = 16000 Hz
    # bits per sample = 16
    # The first is quantization in time
    # The second is quantization in amplitude
    # We also do this for images!
    # 2^16 = 65536 is how many different sound levels we have
    signal = spf.readframes(-1)
    signal = np.fromstring(signal, 'Int16')
    T = len(signal)
    signal = (signal - signal.mean()) / signal.std()

    # signal needs to be of shape N x T(n) x D
    signals = signal.reshape(1, T, 1)

    hmm = HMM(3, 3, 1)
    hmm.init_random(signals)

    init = tf.global_variables_initializer()
    session = tf.InteractiveSession()
    session.run(init)
    hmm.set_session(session)

    hmm.fit(signals, max_iter=30)
//  [python]
def fake_signal():
    signals = get_signals()
    # assume all signals are same length
    signals = np.array(signals)
    hmm = HMM(5, 3, signals[0].shape[1]) # M, K, D
    hmm.init_random(signals)

    init = tf.initialize_all_variables()
    session = tf.InteractiveSession()
    session.run(init)
    hmm.set_session(session)

    hmm.fit(signals, max_iter=30)
    L = hmm.get_cost_multi(signals).sum()
    print("LL for fitted params:", L)

    # test in actual params
    _, _, _, pi, A, R, mu, sigma = big_init()

    # turn these into their "pre-softmax" forms
    pi = np.log(pi)
    A = np.log(A)
    R = np.log(R)
    M, K, D, _ = sigma.shape # need to convert full cov into diag cov
    logSigma = np.zeros((M, K, D))
    for j in range(M):
        for k in range(D):
            logSigma[j,k] = np.log(np.diag(sigma[j,k]))

    hmm.set(pi, A, R, mu, logSigma)
    L = hmm.get_cost_multi(signals).sum()
    print("LL for actual params:", L)
//  [python]
def fake_signal():
    signals = get_signals()
    # assume all signals are same length
    signals = np.array(signals)
    hmm = HMM(5, 3, signals[0].shape[1]) # M, K, D
    hmm.init_random(signals)

    init = tf.global_variables_initializer()
    session = tf.InteractiveSession()
    session.run(init)
    hmm.set_session(session)

    hmm.fit(signals, max_iter=30)
    L = hmm.get_cost_multi(signals).sum()
    print("LL for fitted params:", L)

    # test in actual params
    _, _, _, pi, A, R, mu, sigma = big_init()

    # turn these into their "pre-softmax" forms
    pi = np.log(pi)
    A = np.log(A)
    R = np.log(R)
    M, K, D, _ = sigma.shape # need to convert full cov into diag cov
    logSigma = np.zeros((M, K, D))
    for j in range(M):
        for k in range(D):
            logSigma[j,k] = np.log(np.diag(sigma[j,k]))

    hmm.set(pi, A, R, mu, logSigma)
    L = hmm.get_cost_multi(signals).sum()
    print("LL for actual params:", L)
//  [python]
def testUnknownBatchSizeMobileModel(self):
    batch_size = 1
    height, width = 224, 224
    num_classes = 1000
    with self.test_session() as sess:
      inputs = tf.placeholder(tf.float32, (None, height, width, 3))
      with slim.arg_scope(nasnet.nasnet_mobile_arg_scope()):
        logits, _ = nasnet.build_nasnet_mobile(inputs, num_classes)
      self.assertListEqual(logits.get_shape().as_list(),
                           [None, num_classes])
      images = tf.random_uniform((batch_size, height, width, 3))
      sess.run(tf.initialize_all_variables())
      output = sess.run(logits, {inputs: images.eval()})
      self.assertEquals(output.shape, (batch_size, num_classes))
//  [python]
def testUnknownBatchSizeMobileModel(self):
    batch_size = 1
    height, width = 224, 224
    num_classes = 1000
    with self.test_session() as sess:
      inputs = tf.placeholder(tf.float32, (None, height, width, 3))
      with slim.arg_scope(nasnet.nasnet_mobile_arg_scope()):
        logits, _ = nasnet.build_nasnet_mobile(inputs, num_classes)
      self.assertListEqual(logits.get_shape().as_list(),
                           [None, num_classes])
      images = tf.random_uniform((batch_size, height, width, 3))
      sess.run(tf.global_variables_initializer())
      output = sess.run(logits, {inputs: images.eval()})
      self.assertEquals(output.shape, (batch_size, num_classes))
//  [python]
def testEvaluationMobileModel(self):
    batch_size = 2
    height, width = 224, 224
    num_classes = 1000
    with self.test_session() as sess:
      eval_inputs = tf.random_uniform((batch_size, height, width, 3))
      with slim.arg_scope(nasnet.nasnet_mobile_arg_scope()):
        logits, _ = nasnet.build_nasnet_mobile(eval_inputs,
                                               num_classes,
                                               is_training=False)
      predictions = tf.argmax(logits, 1)
      sess.run(tf.initialize_all_variables())
      output = sess.run(predictions)
      self.assertEquals(output.shape, (batch_size,))
//  [python]
def testEvaluationMobileModel(self):
    batch_size = 2
    height, width = 224, 224
    num_classes = 1000
    with self.test_session() as sess:
      eval_inputs = tf.random_uniform((batch_size, height, width, 3))
      with slim.arg_scope(nasnet.nasnet_mobile_arg_scope()):
        logits, _ = nasnet.build_nasnet_mobile(eval_inputs,
                                               num_classes,
                                               is_training=False)
      predictions = tf.argmax(logits, 1)
      sess.run(tf.global_variables_initializer())
      output = sess.run(predictions)
      self.assertEquals(output.shape, (batch_size,))
//  [python]
def predict(track, model_config, load_model, results_dir=None):
    '''
    Function in accordance with MUSB evaluation API. Takes MUSDB track object and computes corresponding source estimates, as well as calls evlauation script.
    Model has to be saved beforehand into a pickle file containing model configuration dictionary and checkpoint path!
    :param track: Track object
    :param results_dir: Directory where SDR etc. values should be saved
    :return: Source estimates dictionary
    '''

    # Determine input and output shapes, if we use U-net as separator
    disc_input_shape = [model_config["batch_size"], model_config["num_frames"], 0]  # Shape of discriminator input
    if model_config["network"] == "unet":
        separator_class = Models.UnetAudioSeparator.UnetAudioSeparator(model_config)
    elif model_config["network"] == "unet_spectrogram":
        separator_class = Models.UnetSpectrogramSeparator.UnetSpectrogramSeparator(model_config)
    else:
        raise NotImplementedError

    sep_input_shape, sep_output_shape = separator_class.get_padding(np.array(disc_input_shape))
    separator_func = separator_class.get_output

    # Batch size of 1
    sep_input_shape[0] = 1
    sep_output_shape[0] = 1

    mix_ph = tf.placeholder(tf.float32, sep_input_shape)

    print("Testing...")

    # BUILD MODELS
    # Separator
    separator_sources = separator_func(mix_ph, training=False, return_spectrogram=False, reuse=False)

    # Start session and queue input threads
    sess = tf.Session()
    sess.run(tf.initialize_all_variables())

    # Load model
    # Load pretrained model to continue training, if we are supposed to
    restorer = tf.train.Saver(None, write_version=tf.train.SaverDef.V2)
    print("Num of variables" + str(len(tf.global_variables())))
    restorer.restore(sess, load_model)
    print('Pre-trained model restored for song prediction')

    mix_audio, orig_sr, mix_channels = track.audio, track.rate, track.audio.shape[1] # Audio has (n_samples, n_channels) shape
    separator_preds = predict_track(model_config, sess, mix_audio, orig_sr, sep_input_shape, sep_output_shape, separator_sources, mix_ph)

    # Upsample predicted source audio and convert to stereo. Make sure to resample back to the exact number of samples in the original input (with fractional orig_sr/new_sr this causes issues otherwise)
    pred_audio = {name : Utils.resample(separator_preds[name], model_config["expected_sr"], orig_sr)[:mix_audio.shape[0],:] for name in model_config["source_names"]}

    if model_config["mono_downmix"] and mix_channels > 1: # Convert to multichannel if mixture input was multichannel by duplicating mono estimate
        pred_audio = {name : np.tile(pred_audio[name], [1, mix_channels]) for name in list(pred_audio.keys())}

    # Evaluate using museval, if we are currently evaluating MUSDB
    if results_dir is not None:
        scores = museval.eval_mus_track(track, pred_audio, output_dir=results_dir)

        # print nicely formatted mean scores
        print(scores)

    # Close session, clear computational graph
    sess.close()
    tf.reset_default_graph()

    return pred_audio
//  [python]
def predict(track, model_config, load_model, results_dir=None):
    '''
    Function in accordance with MUSB evaluation API. Takes MUSDB track object and computes corresponding source estimates, as well as calls evlauation script.
    Model has to be saved beforehand into a pickle file containing model configuration dictionary and checkpoint path!
    :param track: Track object
    :param results_dir: Directory where SDR etc. values should be saved
    :return: Source estimates dictionary
    '''

    # Determine input and output shapes, if we use U-net as separator
    disc_input_shape = [model_config["batch_size"], model_config["num_frames"], 0]  # Shape of discriminator input
    if model_config["network"] == "unet":
        separator_class = Models.UnetAudioSeparator.UnetAudioSeparator(model_config)
    elif model_config["network"] == "unet_spectrogram":
        separator_class = Models.UnetSpectrogramSeparator.UnetSpectrogramSeparator(model_config)
    else:
        raise NotImplementedError

    sep_input_shape, sep_output_shape = separator_class.get_padding(np.array(disc_input_shape))
    separator_func = separator_class.get_output

    # Batch size of 1
    sep_input_shape[0] = 1
    sep_output_shape[0] = 1

    mix_ph = tf.placeholder(tf.float32, sep_input_shape)

    print("Testing...")

    # BUILD MODELS
    # Separator
    separator_sources = separator_func(mix_ph, training=False, return_spectrogram=False, reuse=False)

    # Start session and queue input threads
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    # Load model
    # Load pretrained model to continue training, if we are supposed to
    restorer = tf.train.Saver(None, write_version=tf.train.SaverDef.V2)
    print("Num of variables" + str(len(tf.global_variables())))
    restorer.restore(sess, load_model)
    print('Pre-trained model restored for song prediction')

    mix_audio, orig_sr, mix_channels = track.audio, track.rate, track.audio.shape[1] # Audio has (n_samples, n_channels) shape
    separator_preds = predict_track(model_config, sess, mix_audio, orig_sr, sep_input_shape, sep_output_shape, separator_sources, mix_ph)

    # Upsample predicted source audio and convert to stereo. Make sure to resample back to the exact number of samples in the original input (with fractional orig_sr/new_sr this causes issues otherwise)
    pred_audio = {name : Utils.resample(separator_preds[name], model_config["expected_sr"], orig_sr)[:mix_audio.shape[0],:] for name in model_config["source_names"]}

    if model_config["mono_downmix"] and mix_channels > 1: # Convert to multichannel if mixture input was multichannel by duplicating mono estimate
        pred_audio = {name : np.tile(pred_audio[name], [1, mix_channels]) for name in list(pred_audio.keys())}

    # Evaluate using museval, if we are currently evaluating MUSDB
    if results_dir is not None:
        scores = museval.eval_mus_track(track, pred_audio, output_dir=results_dir)

        # print nicely formatted mean scores
        print(scores)

    # Close session, clear computational graph
    sess.close()
    tf.reset_default_graph()

    return pred_audio
//  [python]
def init(self):
        """Call .init() before training starts.
        - Initialize the variables.
        - Save the model config into json file.
        """
        self.sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])
        self._is_init = True
        self.step = 0

        self.save_checkpoint()  # make sure saver is created.
        # Save the model config into a json.
        config_path = os.path.join(self.checkpoint_dir, 'model.config.json')
        with open(config_path, 'w') as fout:
            json.dump(self.config, fout)
//  [python]
def createModel(session, forward_only,from_vocab_size,to_vocab_size):
    """Create translation model and initialize or load parameters in session."""
    model = seq2seq_model.Seq2SeqModel(
      from_vocab_size,#from
      to_vocab_size,#to
      _buckets,
      hidden_size,
      num_layers,
      dropout,
      grad_clip,
      batch_size,
      learning_rate,
      lr_decay_factor,
      forward_only=forward_only,
      dtype=tf.float32)
      
    print("model is ok")

    
    ckpt = tf.train.latest_checkpoint(checkpoint_dir)
    if ckpt!=None:
        model.saver.restore(session, ckpt)
        print ("Reading model parameters from {0}".format(ckpt))
    else:
        print ("Created model with fresh parameters.")
        session.run(tf.initialize_all_variables())  

    return model
//  [python]
def createModel(session, forward_only,from_vocab_size,to_vocab_size):
    """Create translation model and initialize or load parameters in session."""
    model = seq2seq_model.Seq2SeqModel(
      from_vocab_size,#from
      to_vocab_size,#to
      _buckets,
      hidden_size,
      num_layers,
      dropout,
      grad_clip,
      batch_size,
      learning_rate,
      lr_decay_factor,
      forward_only=forward_only,
      dtype=tf.float32)
      
    print("model is ok")

    
    ckpt = tf.train.latest_checkpoint(checkpoint_dir)
    if ckpt!=None:
        model.saver.restore(session, ckpt)
        print ("Reading model parameters from {0}".format(ckpt))
    else:
        print ("Created model with fresh parameters.")
        session.run(tf.global_variables_initializer())  

    return model
//  [python]
def setup(self, scheme, groups, preprocess):
        self.new_batch = partial(
            EpisodeBatchNP,
            scheme,
            groups,
            1,  # Note: batch size must be 1 in a episode
            self.episode_limit + 1,
            preprocess=preprocess,
        )
        self.alg.build_actor_graph()  # 1 only use for explore !

        self.alg.build_train_graph()

        # note: init with only once are importance!
        with self.alg.graph.as_default():
            self.alg.sess.run(tf.global_variables_initializer())

            self.alg.hi_out_val_default = self.alg.sess.run(
                self.alg.gru_cell.zero_state(self.args.n_agents, dtype=tf.float32))

        writer = tf.summary.FileWriter(logdir="logdir", graph=self.alg.graph)
        writer.flush()
//  [python]
def restore(self, session,model_name="model.max.ckpt"):
    # Don't try to restore unused variables from the TF-Hub ELMo module.
    if self.with_split_antecedent:
      vars_to_restore = [v for v in tf.global_variables() if "module/" not in v.name]
    else:
      vars_to_restore = [v for v in tf.global_variables() if "module/" not in v.name and "plural_scoring/" not in v.name]
    saver = tf.train.Saver(vars_to_restore)
    checkpoint_path = os.path.join(self.config["log_dir"], model_name)
    print("Restoring from {}".format(checkpoint_path))
    session.run(tf.initialize_all_variables())
    saver.restore(session, checkpoint_path)
//  [python]
def restore(self, session,model_name="model.max.ckpt"):
    # Don't try to restore unused variables from the TF-Hub ELMo module.
    if self.with_split_antecedent:
      vars_to_restore = [v for v in tf.global_variables() if "module/" not in v.name]
    else:
      vars_to_restore = [v for v in tf.global_variables() if "module/" not in v.name and "plural_scoring/" not in v.name]
    saver = tf.train.Saver(vars_to_restore)
    checkpoint_path = os.path.join(self.config["log_dir"], model_name)
    print("Restoring from {}".format(checkpoint_path))
    session.run(tf.global_variables_initializer())
    saver.restore(session, checkpoint_path)
//  [python]
def _run(self, imgs_l, imgs_ab, imgs_true_ab, opt_operations):
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
            sess.run(tf.local_variables_initializer())

            # Coordinate the loading of image files.
            coord = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(coord=coord)

            res = sess.run(
                {"imgs_l": imgs_l, "imgs_ab": imgs_ab, "imgs_true_ab": imgs_true_ab,}
            )

            img_gray = l_to_rgb(res["imgs_l"][0][:, :, 0])
            img_output = lab_to_rgb(res["imgs_l"][0][:, :, 0], res["imgs_ab"][0])
            img_true = lab_to_rgb(res["imgs_l"][0][:, :, 0], res["imgs_true_ab"][0])

            plt.subplot(2, 3, 1)
            plt.imshow(img_gray)
            plt.title("Input (grayscale)")
            plt.axis("off")
            plt.subplot(2, 3, 2)
            plt.imshow(img_output)
            plt.title("Network output")
            plt.axis("off")
            plt.subplot(2, 3, 3)
            plt.imshow(img_true)
            plt.title("Target (original)")
            plt.axis("off")

            for epoch in range(NUM_EPOCHS):
                print("Epoch:", epoch, end=" ")
                res = sess.run(opt_operations)
                print("Cost:", res["cost"])

            res = sess.run(
                {"imgs_l": imgs_l, "imgs_ab": imgs_ab, "imgs_true_ab": imgs_true_ab,}
            )

            img_gray = l_to_rgb(res["imgs_l"][0][:, :, 0])
            img_output = lab_to_rgb(res["imgs_l"][0][:, :, 0], res["imgs_ab"][0])
            img_true = lab_to_rgb(res["imgs_l"][0][:, :, 0], res["imgs_true_ab"][0])

            plt.subplot(2, 3, 4)
            plt.imshow(img_gray)
            plt.title("Input (grayscale)")
            plt.axis("off")
            plt.subplot(2, 3, 5)
            plt.imshow(img_output)
            plt.title("Network output")
            plt.axis("off")
            plt.subplot(2, 3, 6)
            plt.imshow(img_true)
            plt.title("Target (original)")
            plt.axis("off")

            plt.show()

            # Finish off the queue coordinator.
            coord.request_stop()
            coord.join(threads)
//  [python]
def create_or_load_model(model, model_dir, session, name):
    """Create translation model and initialize or load parameters in session."""
    latest_ckpt = tf.train.latest_checkpoint(model_dir)
    if latest_ckpt:
        model._replace(model=load_model(model.model, latest_ckpt, session, name))
        utils.print_out("checkpoint found, load checkpoint\n %s" % latest_ckpt)
    else:
        utils.print_out("  checkpoint not found in %s" % (model_dir))
        utils.print_out("  created %s model with fresh parameters" % (name))
        session.run(tf.initialize_all_variables())
        # session.run(tf.tables_initializer())
    global_step = model.model.global_step.eval(session=session)
    epoch_num = model.model.epoch_num.eval(session=session)
    return model, global_step, epoch_num
//  [python]
def create_or_load_model(model, model_dir, session, name):
    """Create translation model and initialize or load parameters in session."""
    latest_ckpt = tf.train.latest_checkpoint(model_dir)
    if latest_ckpt:
        model._replace(model=load_model(model.model, latest_ckpt, session, name))
        utils.print_out("checkpoint found, load checkpoint\n %s" % latest_ckpt)
    else:
        utils.print_out("  checkpoint not found in %s" % (model_dir))
        utils.print_out("  created %s model with fresh parameters" % (name))
        session.run(tf.global_variables_initializer())
        # session.run(tf.tables_initializer())
    global_step = model.model.global_step.eval(session=session)
    epoch_num = model.model.epoch_num.eval(session=session)
    return model, global_step, epoch_num
//  [python]
def predict():
    num_classes = 15
    learning_rate = 0.001
    batch_size = 8
    decay_steps = 1000
    decay_rate = 0.9
    sequence_length = 10
    vocab_size = 10000
    embed_size = 100
    hidden_size = 100
    is_training = False
    story_length = 3
    dropout_keep_prob = 1
    model = EntityNetwork(num_classes, learning_rate, batch_size, decay_steps, decay_rate, sequence_length,
                          story_length, vocab_size, embed_size, hidden_size, is_training,
                          multi_label_flag=False, block_size=20)
    ckpt_dir = 'checkpoint_entity_network/dummy_test/'
    saver = tf.train.Saver()
    with tf.Session() as sess:
        sess.run(tf.initialize_all_variables())
        saver.restore(sess, tf.train.latest_checkpoint(ckpt_dir))
        for i in range(1500):
            story = np.random.randn(batch_size, story_length, sequence_length)
            story[story > 0] = 1
            story[story <= 0] = 0
            query = np.random.randn(batch_size, sequence_length)  # [batch_size, sequence_length]
            query[query > 0] = 1
            query[query <= 0] = 0
            answer_single = np.sum(query, axis=1) + np.round(0.1 * np.sum(np.sum(story, axis=1),
                                                                          axis=1))  # [batch_size].e.g. np.array([1, 0, 1, 1, 1, 2, 1, 1])
            predict = sess.run([model.predictions], feed_dict={model.query: query, model.story: story,
                                                               model.dropout_keep_prob: dropout_keep_prob})
            print(i, "query:", query, "=====================>")
            print(i, "label:", answer_single, "prediction:", predict)
//  [python]
def predict():
    num_classes = 15
    learning_rate = 0.001
    batch_size = 8
    decay_steps = 1000
    decay_rate = 0.9
    sequence_length = 10
    vocab_size = 10000
    embed_size = 100
    hidden_size = 100
    is_training = False
    story_length = 3
    dropout_keep_prob = 1
    model = EntityNetwork(num_classes, learning_rate, batch_size, decay_steps, decay_rate, sequence_length,
                          story_length, vocab_size, embed_size, hidden_size, is_training,
                          multi_label_flag=False, block_size=20)
    ckpt_dir = 'checkpoint_entity_network/dummy_test/'
    saver = tf.train.Saver()
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        saver.restore(sess, tf.train.latest_checkpoint(ckpt_dir))
        for i in range(1500):
            story = np.random.randn(batch_size, story_length, sequence_length)
            story[story > 0] = 1
            story[story <= 0] = 0
            query = np.random.randn(batch_size, sequence_length)  # [batch_size, sequence_length]
            query[query > 0] = 1
            query[query <= 0] = 0
            answer_single = np.sum(query, axis=1) + np.round(0.1 * np.sum(np.sum(story, axis=1),
                                                                          axis=1))  # [batch_size].e.g. np.array([1, 0, 1, 1, 1, 2, 1, 1])
            predict = sess.run([model.predictions], feed_dict={model.query: query, model.story: story,
                                                               model.dropout_keep_prob: dropout_keep_prob})
            print(i, "query:", query, "=====================>")
            print(i, "label:", answer_single, "prediction:", predict)
//  [python]
def __init__(self, model, restore_snap=None):

        # Add a softmax operation for predictions
        model.prob_logits = tf.nn.softmax(model.logits)

        # Tensorflow Saver definition
        my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='KernelPointNetwork')
        self.saver = tf.train.Saver(my_vars, max_to_keep=100)

        # Create a session for running Ops on the Graph.
        on_CPU = False
        if on_CPU:
            cProto = tf.ConfigProto(device_count={'GPU': 0})
        else:
            cProto = tf.ConfigProto()
            cProto.gpu_options.allow_growth = True
        self.sess = tf.Session(config=cProto)

        # Init variables
        self.sess.run(tf.initialize_all_variables())

        # Name of the snapshot to restore to (None if you want to start from beginning)
        # restore_snap = join(self.saving_path, 'snapshots/snap-40000')
        if (restore_snap is not None):
            self.saver.restore(self.sess, restore_snap)
            print("Model restored.")
//  [python]
def __init__(self, model, restore_snap=None):

        # Add a softmax operation for predictions
        model.prob_logits = tf.nn.softmax(model.logits)

        # Tensorflow Saver definition
        my_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='KernelPointNetwork')
        self.saver = tf.train.Saver(my_vars, max_to_keep=100)

        # Create a session for running Ops on the Graph.
        on_CPU = False
        if on_CPU:
            cProto = tf.ConfigProto(device_count={'GPU': 0})
        else:
            cProto = tf.ConfigProto()
            cProto.gpu_options.allow_growth = True
        self.sess = tf.Session(config=cProto)

        # Init variables
        self.sess.run(tf.global_variables_initializer())

        # Name of the snapshot to restore to (None if you want to start from beginning)
        # restore_snap = join(self.saving_path, 'snapshots/snap-40000')
        if (restore_snap is not None):
            self.saver.restore(self.sess, restore_snap)
            print("Model restored.")
//  [python]
def _init_vars(self):
        self.sess.run(tf.initialize_all_variables())
//  [python]
def _init_vars(self):
        self.sess.run(tf.global_variables_initializer())
//  [python]
def test():
    # below is a function test; if you use this for text classifiction, you need to tranform sentence to indices of vocabulary first. then feed data to the graph.
    num_classes = 15
    learning_rate = 0.001
    batch_size = 8
    decay_steps = 1000
    decay_rate = 0.9
    sequence_length = 10
    vocab_size = 10000
    embed_size = 100
    hidden_size = 100
    is_training = True
    story_length = 3
    dropout_keep_prob = 1
    use_bi_lstm=False
    model = EntityNetwork(num_classes, learning_rate, batch_size, decay_steps, decay_rate, sequence_length,
                          story_length, vocab_size, embed_size, hidden_size, is_training,
                          multi_label_flag=False, block_size=20,use_bi_lstm=use_bi_lstm)
    ckpt_dir = 'checkpoint_entity_network/dummy_test/'
    saver = tf.train.Saver()
    with tf.Session() as sess:
        sess.run(tf.initialize_all_variables())
        for i in range(1500):
            # input_x should be:[batch_size, num_sentences,self.sequence_length]
            story = np.random.randn(batch_size, story_length, sequence_length)
            story[story > 0] = 1
            story[story <= 0] = 0
            query = np.random.randn(batch_size, sequence_length)  # [batch_size, sequence_length]
            query[query > 0] = 1
            query[query <= 0] = 0
            answer_single = np.sum(query, axis=1) + np.round(0.1 * np.sum(np.sum(story, axis=1),
                                                                          axis=1))  # [batch_size].e.g. np.array([1, 0, 1, 1, 1, 2, 1, 1])
            loss, acc, predict, _ = sess.run(
                [model.loss_val, model.accuracy, model.predictions, model.train_op],
                feed_dict={model.query: query, model.story: story, model.answer_single: answer_single,
                           model.dropout_keep_prob: dropout_keep_prob})
            print(i, "query:", query, "=====================>")
            print(i, "loss:", loss, "acc:", acc, "label:", answer_single, "prediction:", predict)
            if i % 300 == 0:
                save_path = ckpt_dir + "model.ckpt"
                saver.save(sess, save_path, global_step=i * 300)
//  [python]
def test():
    # below is a function test; if you use this for text classifiction, you need to tranform sentence to indices of vocabulary first. then feed data to the graph.
    num_classes = 15
    learning_rate = 0.001
    batch_size = 8
    decay_steps = 1000
    decay_rate = 0.9
    sequence_length = 10
    vocab_size = 10000
    embed_size = 100
    hidden_size = 100
    is_training = True
    story_length = 3
    dropout_keep_prob = 1
    use_bi_lstm=False
    model = EntityNetwork(num_classes, learning_rate, batch_size, decay_steps, decay_rate, sequence_length,
                          story_length, vocab_size, embed_size, hidden_size, is_training,
                          multi_label_flag=False, block_size=20,use_bi_lstm=use_bi_lstm)
    ckpt_dir = 'checkpoint_entity_network/dummy_test/'
    saver = tf.train.Saver()
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        for i in range(1500):
            # input_x should be:[batch_size, num_sentences,self.sequence_length]
            story = np.random.randn(batch_size, story_length, sequence_length)
            story[story > 0] = 1
            story[story <= 0] = 0
            query = np.random.randn(batch_size, sequence_length)  # [batch_size, sequence_length]
            query[query > 0] = 1
            query[query <= 0] = 0
            answer_single = np.sum(query, axis=1) + np.round(0.1 * np.sum(np.sum(story, axis=1),
                                                                          axis=1))  # [batch_size].e.g. np.array([1, 0, 1, 1, 1, 2, 1, 1])
            loss, acc, predict, _ = sess.run(
                [model.loss_val, model.accuracy, model.predictions, model.train_op],
                feed_dict={model.query: query, model.story: story, model.answer_single: answer_single,
                           model.dropout_keep_prob: dropout_keep_prob})
            print(i, "query:", query, "=====================>")
            print(i, "loss:", loss, "acc:", acc, "label:", answer_single, "prediction:", predict)
            if i % 300 == 0:
                save_path = ckpt_dir + "model.ckpt"
                saver.save(sess, save_path, global_step=i * 300)