Python tensorflow.arg_max() Examples

def test_batch_iou(self):
    with self.test_session() as sess:
      anchors = set_anchors(img_shape=[config.IMG_HEIGHT, config.IMG_WIDTH],
                            fea_shape=[config.FEA_HEIGHT, config.FEA_WIDTH])
      anchors_shape = anchors.get_shape().as_list()
      fea_h = anchors_shape[0]
      fea_w = anchors_shape[1]
      num_anchors = anchors_shape[2] * fea_h * fea_w
      anchors = tf.reshape(anchors, [num_anchors, 4])  # reshape anchors
      anchors = xywh_to_yxyx(anchors)
      bbox = tf.constant([0.75, 0.75, 0.2, 0.2], dtype=tf.float32)
      bbox = xywh_to_yxyx(bbox)
      iou = batch_iou(anchors, bbox)
      anchor_idx = tf.arg_max(iou, dimension=0)
      anchors, output, anchor_idx = sess.run([anchors, iou, anchor_idx])
      print(anchors)
      print(output)
      print(anchor_idx) 

def __call__(self, sess, epoch, iteration, model, loss, processed):
        if iteration == 0 and epoch % self.at_every_epoch == 0:
            total = 0
            correct = 0
            for values in self.batcher:
                total += len(values[-1])
                feed_dict = {}
                for i in range(0, len(self.placeholders)):
                    feed_dict[self.placeholders[i]] = values[i]
                truth = np.argmax(values[-1], 1)
                predicted = sess.run(tf.arg_max(tf.nn.softmax(model), 1),
                                     feed_dict=feed_dict)
                correct += sum(truth == predicted)
            acc = float(correct) / total
            self.update_summary(sess, iteration, ACCURACY_TRACE_TAG, acc)
            print("Epoch " + str(epoch) +
                  "\tAcc " + str(acc) +
                  "\tCorrect " + str(correct) + "\tTotal " + str(total)) 

def batch_iou_(anchors, bboxes):
  """ Compute iou of two batch of boxes. Box format '[y_min, x_min, y_max, x_max]'.
  Args:
    anchors: know shape
    bboxes: dynamic shape
  Return:
    ious: 2-D with shape '[num_bboxes, num_anchors]'
    indices: [num_bboxes, 1]
  """
  num_anchors = anchors.get_shape().as_list()[0]
  ious_list = []
  for i in range(num_anchors):
    anchor = anchors[i]
    _ious = batch_iou(bboxes, anchor)
    ious_list.append(_ious)
  ious = tf.stack(ious_list, axis=0)
  ious = tf.transpose(ious)

  indices = tf.arg_max(ious, dimension=1)

  return ious, indices 

def initialize_neural_network_model(self, X_train):
        # Create placeholders
        self.features_pl = tf.placeholder(tf.float32, [None, len(X_train[0])], 'features')
        self.stances_pl = tf.placeholder(tf.int64, [None], 'stances')
        self.keep_prob_pl = tf.placeholder(tf.float32)

        # Infer batch size
        self.batch_size = tf.shape(self.features_pl)[0]

        # Define multi-layer perceptron
        self.hidden_layer = tf.nn.dropout(tf.nn.relu(tf.contrib.layers.linear(self.features_pl, self.hidden_size)), keep_prob=self.keep_prob_pl)
        self.logits_flat = tf.nn.dropout(tf.contrib.layers.linear(self.hidden_layer, self.target_size), keep_prob=self.keep_prob_pl)
        self.logits = tf.reshape(self.logits_flat, [self.batch_size, self.target_size])

        # Define L2 loss
        self.tf_vars = tf.trainable_variables()
        self.l2_loss = tf.add_n([tf.nn.l2_loss(v) for v in self.tf_vars if 'bias' not in v.name]) * self.l2_alpha

        # Define overall loss
        self.loss = tf.reduce_sum(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.logits, labels=self.stances_pl) + self.l2_loss)

        # Define prediction
        self.softmaxed_logits = tf.nn.softmax(self.logits)
        self.predict_value = tf.arg_max(self.softmaxed_logits, 1) 

def build_output(self, inputs, inferences):
    scores = tf.nn.softmax(inferences, name='scores')
    tf.add_to_collection('outputs', scores)

    with tf.name_scope('labels'):
      label_indices = tf.arg_max(inferences, 1, name='arg_max')
      labels = self.classification.output_labels(label_indices)
      tf.add_to_collection('outputs', labels)

    keys = self.classification.keys(inputs)
    if keys:
      # Key feature, if it exists, is a passthrough to the output.
      # The use of identity is to name the tensor and correspondingly the output field.
      keys = tf.identity(keys, name='key')
      tf.add_to_collection('outputs', keys)

    return {
      'label': labels,
      'score': scores
    } 

def max_sentence_similarity(sentence_input, similarity_matrix):
    """
    Parameters
    ----------
    sentence_input: Tensor
        Tensor of shape (batch_size, num_sentence_words, rnn_hidden_dim).

    similarity_matrix: Tensor
        Tensor of shape (batch_size, num_sentence_words, num_sentence_words).
    """
    # Shape: (batch_size, passage_len)
    def single_instance(inputs):
        single_sentence = inputs[0]
        argmax_index = inputs[1]
        # Shape: (num_sentence_words, rnn_hidden_dim)
        return tf.gather(single_sentence, argmax_index)

    question_index = tf.arg_max(similarity_matrix, 2)
    elems = (sentence_input, question_index)
    # Shape: (batch_size, num_sentence_words, rnn_hidden_dim)
    return tf.map_fn(single_instance, elems, dtype="float") 

def predictor(inputs, targets, target_size):
    init = tf.contrib.layers.xavier_initializer(uniform=True) #uniform=False for truncated normal
    logits = tf.contrib.layers.fully_connected(inputs, target_size, weights_initializer=init, activation_fn=None)

    loss = tf.reduce_mean(
        tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
            labels=targets), name='predictor_loss')
    predict = tf.arg_max(tf.nn.softmax(logits), 1, name='prediction')
    return [logits, loss, predict] 

def predictor(inputs, targets, target_size):
    init = tf.contrib.layers.xavier_initializer(uniform=True) #uniform=False for truncated normal
    logits = tf.contrib.layers.fully_connected(inputs, target_size, weights_initializer=init, activation_fn=None)

    loss = tf.reduce_mean(
        tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
            labels=targets), name='predictor_loss')
    predict = tf.arg_max(tf.nn.softmax(logits), 1, name='prediction')
    return [logits, loss, predict] 

def compute_accuracy(v_xs,v_ys):
    global prediction
    y_pre = sess.run(prediction,feed_dict={xs:v_xs})
    correct_prediction = tf.equal(tf.arg_max(y_pre,1),tf.arg_max(v_ys,1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float32))
    result = sess.run(accuracy,feed_dict={xs:v_xs,ys:v_ys})
    return result 

def iou_from_logits(logits, labels):
  """
  Computes the intersection over union (IoU) score for given logit tensor and target labels
  :param logits: 4D tensor of shape [batch_size, height, width, num_classes]
  :param labels: 3D tensor of shape [batch_size, height, width] and type int32 or int64
  :return: 1D tensor of shape [num_classes] with intersection over union for each class, averaged over batch
  """

  with tf.variable_scope("IoU"):
    # compute predictions
    preds = tf.arg_max(logits, dimension=3)

    num_labels = logits.get_shape().as_list()[-1]
    IoUs = []
    for label in range(num_labels):
      # find pixels with given label
      P = tf.equal(preds, label)
      L = tf.equal(labels, label)

      # Union
      U = tf.logical_or(P, L)
      U = tf.reduce_sum(tf.cast(U, tf.float32))

      # intersection
      I = tf.logical_and(P, L)
      I = tf.reduce_sum(tf.cast(I, tf.float32))

      IoUs.append(I / U)

    return tf.reshape(tf.stack(IoUs), (num_labels,)) 

def compute_iou_from_logits(preds, labels, num_labels):
  """
  Computes the intersection over union (IoU) score for given logit tensor and target labels
  :param logits: 4D tensor of shape [batch_size, height, width, num_classes]
  :param labels: 3D tensor of shape [batch_size, height, width] and type int32 or int64
  :param num_labels: tensor with the number of labels
  :return: 1D tensor of shape [num_classes] with intersection over union for each class, averaged over batch
  """
  with tf.variable_scope("IoU"):
    # compute predictions
    # probs = softmax(logits, axis=-1)
    # preds = tf.arg_max(probs, dimension=3)
    # num_labels = preds.get_shape().as_list()[-1];
    IoUs = []
    for label in range(num_labels):
      # find pixels with given label
      P = tf.equal(preds, label)
      L = tf.equal(labels, label)
      # Union
      U = tf.logical_or(P, L)
      U = tf.reduce_sum(tf.cast(U, tf.float32))
      # intersection
      I = tf.logical_and(P, L)
      I = tf.reduce_sum(tf.cast(I, tf.float32))

      IOU = tf.cast(I, tf.float32) / tf.cast(U, tf.float32)
      # U might be 0!
      IOU = tf.where(tf.equal(U, 0), 1, IOU)
      IOU = tf.Print(IOU, [IOU], "iou" + repr(label))
      IoUs.append(IOU)
    return tf.reshape(tf.stack(IoUs), (num_labels,)) 

def predict(images, exp_config):
    '''
    Returns the prediction for an image given a network from the model zoo
    :param images: An input image tensor
    :param inference_handle: A model function from the model zoo
    :return: A prediction mask, and the corresponding softmax output
    '''

    logits = exp_config.model_handle(images, training=tf.constant(False, dtype=tf.bool), nlabels=exp_config.nlabels)
    softmax = tf.nn.softmax(logits)
    mask = tf.arg_max(softmax, dimension=-1)

    return mask, softmax 

def calc_reward(outputs):
    outputs = outputs[-1] # look at ONLY THE END of the sequence
    outputs = tf.reshape(outputs, (batch_size, cell_out_size))
    h_a_out = weight_variable((cell_out_size, n_classes))

    p_y = tf.nn.softmax(tf.matmul(outputs, h_a_out))
    max_p_y = tf.arg_max(p_y, 1)
    correct_y = tf.cast(labels_placeholder, tf.int64)

    R = tf.cast(tf.equal(max_p_y, correct_y), tf.float32) # reward per example

    reward = tf.reduce_mean(R) # overall reward
    
    p_loc = gaussian_pdf(mean_locs, sampled_locs)
    p_loc = tf.reshape(p_loc, (batch_size, glimpses * 2))

    R = tf.reshape(R, (batch_size, 1))
    J = tf.concat(1, [tf.log(p_y + 1e-5) * onehot_labels_placeholder, tf.log(p_loc + 1e-5) * R])
    J = tf.reduce_sum(J, 1)
    J = tf.reduce_mean(J, 0)
    cost = -J
    
    optimizer = tf.train.AdamOptimizer(lr)
    train_op = optimizer.minimize(cost)

    return cost, reward, max_p_y, correct_y, train_op 

def __init__(self, is_training=True):
        self.graph = tf.Graph()
        with self.graph.as_default():
            if is_training:
                self.x, self.y, self.num_batch = get_batch()
            else: # Evaluation
                self.x = tf.placeholder(tf.int32, shape=(None, hp.max_len,))
                self.y = tf.placeholder(tf.int32, shape=(None, hp.max_len,))
            
            # Character Embedding for x
            self.enc = embed(self.x, len(roma2idx), hp.embed_size, scope="emb_x")
                
            # Encoder
            self.memory = encode(self.enc, is_training=is_training)
            
            # Character Embedding for decoder_inputs
            self.decoder_inputs = shift_by_one(self.y)
            self.dec = embed(self.decoder_inputs, len(surf2idx), hp.embed_size, scope="emb_decoder_inputs")
            
            # Decoder
            self.outputs = decode(self.dec, self.memory, len(surf2idx), is_training=is_training) # (N, T', hp.n_mels*hp.r)
            self.logprobs = tf.log(tf.nn.softmax(self.outputs)+1e-10) 
            self.preds = tf.arg_max(self.outputs, dimension=-1)
                
            if is_training: 
                self.loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=self.y, logits=self.outputs) 
                self.istarget = tf.to_float(tf.not_equal(self.y, tf.zeros_like(self.y))) # masking
                self.mean_loss = tf.reduce_sum(self.loss * self.istarget) / (tf.reduce_sum(self.istarget))
               
                # Training Scheme
                self.global_step = tf.Variable(0, name='global_step', trainable=False)
                self.optimizer = tf.train.AdamOptimizer(learning_rate=hp.lr)
                self.train_op = self.optimizer.minimize(self.mean_loss, global_step=self.global_step)
                   
                # Summary 
                tf.summary.scalar('mean_loss', self.mean_loss)
                self.merged = tf.summary.merge_all() 

def batch_iou_fast(anchors, bboxes):
  """ Compute iou of two batch of boxes. Box format '[y_min, x_min, y_max, x_max]'.
  Args:
    anchors: know shape
    bboxes: dynamic shape
  Return:
    ious: 2-D with shape '[num_bboxes, num_anchors]'
    indices: [num_bboxes, 1]
  """
  num_anchors = anchors.get_shape().as_list()[0]
  tensor_num_bboxes = tf.shape(bboxes)[0]
  indices = tf.reshape(tf.range(tensor_num_bboxes), shape=[-1, 1])
  indices = tf.reshape(tf.stack([indices]*num_anchors, axis=1), shape=[-1, 1])
  bboxes_m = tf.gather_nd(bboxes, indices)

  anchors_m = tf.tile(anchors, [tensor_num_bboxes, 1])

  lr = tf.maximum(
    tf.minimum(bboxes_m[:, 3], anchors_m[:, 3]) -
    tf.maximum(bboxes_m[:, 1], anchors_m[:, 1]),
    0
  )
  tb = tf.maximum(
    tf.minimum(bboxes_m[:, 2], anchors_m[:, 2]) -
    tf.maximum(bboxes_m[:, 0], anchors_m[:, 0]),
    0
  )
  intersection = tf.multiply(tb, lr)
  union = tf.subtract(
    tf.multiply((bboxes_m[:, 3] - bboxes_m[:, 1]), (bboxes_m[:, 2] - bboxes_m[:, 0])) +
    tf.multiply((anchors_m[:, 3] - anchors_m[:, 1]), (anchors_m[:, 2] - anchors_m[:, 0])),
    intersection
  )
  ious = tf.div(intersection, union)

  ious = tf.reshape(ious, shape=[tensor_num_bboxes, num_anchors])

  indices = tf.arg_max(ious, dimension=1)

  return ious, indices 

def __init__(self, is_training=True):
        self.graph = tf.Graph()
        self.is_training=is_training
        with self.graph.as_default():
            if is_training:
                self.x, self.y, self.num_batch = get_batch() 
            else: # Evaluation
                self.x = tf.placeholder(tf.float32, shape=(None, None, hp.n_mels*hp.r))
                self.y = tf.placeholder(tf.int32, shape=(None, hp.max_len))
            
            self.decoder_inputs = embed(shift_by_one(self.y), len(char2idx), hp.embed_size) # (N, T', E)
            
            with tf.variable_scope('net'):
                # Encoder
                self.memory = encode(self.x, is_training=is_training) # (N, T, hp.n_mels*hp.r)
                
                # Decoder
                self.outputs = decode(self.decoder_inputs, self.memory, is_training=is_training) # (N, T', E)
                self.logprobs = tf.log(tf.nn.softmax(self.outputs)+1e-10) 
                self.preds = tf.arg_max(self.outputs, dimension=-1)
                
            if is_training:  
                # Loss
                self.loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=self.y, logits=self.outputs)
                
                # Target masking
                self.istarget = tf.to_float(tf.not_equal(self.y, 0))
                self.mean_loss = tf.reduce_sum(self.loss*self.istarget) / (tf.reduce_sum(self.istarget) + 1e-7)
                
                # Training Scheme
                self.global_step = tf.Variable(0, name='global_step', trainable=False)
                self.optimizer = tf.train.AdamOptimizer(learning_rate=hp.lr)
                self.train_op = self.optimizer.minimize(self.mean_loss, global_step=self.global_step)
                   
                # Summary 
                tf.summary.scalar('mean_loss', self.mean_loss)
                self.merged = tf.summary.merge_all() 

def evaluation(logits, labels, images, nlabels, loss_type):
    '''
    A function for evaluating the performance of the netwrok on a minibatch. This function returns the loss and the 
    current foreground Dice score, and also writes example segmentations and imges to to tensorboard.
    :param logits: Output of network before softmax
    :param labels: Ground-truth label mask
    :param images: Input image mini batch
    :param nlabels: Number of labels in the dataset
    :param loss_type: Which loss should be evaluated
    :return: The loss without weight decay, the foreground dice of a minibatch
    '''

    mask = tf.arg_max(tf.nn.softmax(logits, dim=-1), dimension=-1)  # was 3
    mask_gt = labels

    tf.summary.image('example_gt', prepare_tensor_for_summary(mask_gt, mode='mask', nlabels=nlabels))
    tf.summary.image('example_pred', prepare_tensor_for_summary(mask, mode='mask', nlabels=nlabels))
    tf.summary.image('example_zimg', prepare_tensor_for_summary(images, mode='image'))

    total_loss, nowd_loss, weights_norm = loss(logits, labels, nlabels=nlabels, loss_type=loss_type)

    cdice_structures = losses.per_structure_dice(logits, tf.one_hot(labels, depth=nlabels))
    cdice_foreground = cdice_structures[:,1:]

    cdice = tf.reduce_mean(cdice_foreground)

    return nowd_loss, cdice 

def build_graph(self):
    config = self.config
    self.reader = utils.DataReader(seq_len=config.seq_length, batch_size=config.batch_size, data_filename=config.data_filename)

    self.cell = LayerNormFastWeightsBasicRNNCell(num_units=config.rnn_size)

    self.input_data = tf.placeholder(tf.int32, [None, config.input_length])
    self.targets = tf.placeholder(tf.int32, [None, 1])
    self.initial_state = self.cell.zero_state(tf.shape(self.targets)[0], tf.float32)
    self.initial_fast_weights = self.cell.zero_fast_weights(tf.shape(self.targets)[0], tf.float32)

    with tf.variable_scope("input_embedding"):
      embedding = tf.get_variable("embedding", [config.vocab_size, config.embedding_size])
      inputs = tf.split(1, config.input_length, tf.nn.embedding_lookup(embedding, self.input_data))
      inputs = [tf.squeeze(input, [1]) for input in inputs]

    with tf.variable_scope("send_to_rnn"):
      state = (self.initial_state, self.initial_fast_weights)
      output = None

      for i, input in enumerate(inputs):
        if i > 0:
          tf.get_variable_scope().reuse_variables()
        output, state = self.cell(input, state)

    with tf.variable_scope("softmax"):
      softmax_w = tf.get_variable("softmax_w", [config.rnn_size, config.vocab_size])
      softmax_b = tf.get_variable("softmax_b", [config.vocab_size])
      self.logits = tf.matmul(output, softmax_w) + softmax_b
      self.probs = tf.nn.softmax(self.logits)
      self.output = tf.cast(tf.reshape(tf.arg_max(self.probs, 1), [-1, 1]), tf.int32)
      self.accuracy = tf.reduce_mean(tf.cast(tf.equal(self.output, self.targets), tf.float32))

    loss = seq2seq.sequence_loss_by_example([self.logits],
                                            [tf.reshape(self.targets, [-1])],
                                            [tf.ones([config.batch_size])],
                                            config.vocab_size)

    self.cost = tf.reduce_mean(loss)
    self.final_state = state

    # self.lr = tf.Variable(0.001, trainable=False)
    tvars = tf.trainable_variables()
    grads, _ = tf.clip_by_global_norm(tf.gradients(self.cost, tvars),
                                      config.grad_clip)
    optimizer = tf.train.AdamOptimizer()  # self.lr)
    self.train_op = optimizer.apply_gradients(zip(grads, tvars))

    self.summary_accuracy = tf.scalar_summary('accuracy', self.accuracy)
    tf.scalar_summary('cost', self.cost)
    self.summary_all = tf.merge_all_summaries() 

def build_graph(self):
    config = self.config
    self.reader = utils.DataReader(seq_len=config.seq_length, batch_size=config.batch_size, data_filename=config.data_filename)

    self.cell = rnn_cell.BasicLSTMCell(config.rnn_size, state_is_tuple=True)

    self.input_data = tf.placeholder(tf.int32, [None, config.input_length])
    self.targets = tf.placeholder(tf.int32, [None, 1])
    self.initial_state = self.cell.zero_state(tf.shape(self.targets)[0], tf.float32)

    with tf.variable_scope("input_embedding"):
      embedding = tf.get_variable("embedding", [config.vocab_size, config.rnn_size])
      inputs = tf.split(1, config.input_length, tf.nn.embedding_lookup(embedding, self.input_data))
      inputs = [tf.squeeze(input, [1]) for input in inputs]

    with tf.variable_scope("send_to_rnn"):
      state = self.initial_state
      output = None

      for i, input in enumerate(inputs):
        if i > 0:
          tf.get_variable_scope().reuse_variables()
        output, state = self.cell(input, state)

    with tf.variable_scope("softmax"):
      softmax_w = tf.get_variable("softmax_w", [config.rnn_size, config.vocab_size])
      softmax_b = tf.get_variable("softmax_b", [config.vocab_size])
      self.logits = tf.matmul(output, softmax_w) + softmax_b
      self.probs = tf.nn.softmax(self.logits)
      self.output = tf.cast(tf.reshape(tf.arg_max(self.probs, 1), [-1, 1]), tf.int32)
      self.accuracy = tf.reduce_mean(tf.cast(tf.equal(self.output, self.targets), tf.float32))

    loss = seq2seq.sequence_loss_by_example([self.logits],
                                            [tf.reshape(self.targets, [-1])],
                                            [tf.ones([config.batch_size])],
                                            config.vocab_size)

    self.cost = tf.reduce_mean(loss)
    self.final_state = state

    # self.lr = tf.Variable(0.001, trainable=False)
    tvars = tf.trainable_variables()
    grads, _ = tf.clip_by_global_norm(tf.gradients(self.cost, tvars),
                                      config.grad_clip)
    optimizer = tf.train.AdamOptimizer()#self.lr)
    self.train_op = optimizer.apply_gradients(zip(grads, tvars))

    self.summary_accuracy = tf.scalar_summary('accuracy', self.accuracy)
    tf.scalar_summary('cost', self.cost)
    self.summary_all = tf.merge_all_summaries() 

def __init__(self, sess, config, name, is_train):
    self.sess = sess
    self.name = name
    self.is_train = is_train


    self.X_hsd = tf.placeholder(tf.float32, shape=[config.batch_size, config.im_size, config.im_size, 3], name="original_color_image")
    self.D, h_s = tf.split(self.X_hsd,[1,2], axis=3)

    self.E_Step = CNN("E_Step", config, is_train=self.is_train)
    self.Gama = self.E_Step(self.D)
    self.loss, self.Mu, self.Std = GMM_M_Step(self.X_hsd, self.Gama, config.ClusterNo, name='GMM_Statistics')
    
    if self.is_train:

      self.optim = tf.train.AdamOptimizer(config.lr)
      self.train = self.optim.minimize(self.loss, var_list=self.E_Step.Param)

    ClsLbl = tf.arg_max(self.Gama, 3)
    ClsLbl = tf.cast(ClsLbl, tf.float32)
    
    ColorTable = [[255,0,0],[0,255,0],[0,0,255],[255,255,0], [0,255,255], [255,0,255]]
    colors = tf.cast(tf.constant(ColorTable), tf.float32)
    Msk = tf.tile(tf.expand_dims(ClsLbl, axis=3),[1,1,1,3])
    for k in range(0, config.ClusterNo):
        ClrTmpl = tf.einsum('anmd,df->anmf', tf.expand_dims(tf.ones_like(ClsLbl), axis=3), tf.reshape(colors[k,...],[1,3]))
        Msk = tf.where(tf.equal(Msk,k), ClrTmpl, Msk)
    
    
    self.X_rgb = utils.HSD2RGB(self.X_hsd)
    tf.summary.image("1.Input_image", self.X_rgb*255.0, max_outputs=2)
    tf.summary.image("2.Gamma_image",  Msk, max_outputs=2)
    tf.summary.image("3.Density_image", self.D*255.0, max_outputs=2)
    tf.summary.scalar("loss", self.loss)

    self.summary_op = tf.summary.merge_all()

    self.saver = tf.train.Saver()
    self.summary_writer = tf.summary.FileWriter(config.logs_dir, self.sess.graph)

    self.sess.run(tf.global_variables_initializer())
    
    ckpt = tf.train.get_checkpoint_state(config.logs_dir)
    if ckpt and ckpt.model_checkpoint_path:
        self.saver.restore(self.sess, ckpt.model_checkpoint_path)
        print("Model restored...") 

def test_ensemble(self, M =1):
        """
        Test using the simpliest ensemble methods: max voting
        But this implemetation is not used, because it is complicated, we just test run the same task instance for
        several times and max voting the results. In experiemnt_builder.py.
        :return:
        """
        with tf.name_scope("losses"):
            [b, num_classes, spc] = self.support_set_labels.get_shape().as_list()
            print("data type ", self.support_set_labels.dtype)
            self.support_set_labels_ = tf.reshape(self.support_set_labels, shape=(b, num_classes * spc))
            print("data type ", self.support_set_labels.dtype)
            self.support_set_labels_ = tf.one_hot(self.support_set_labels_, self.num_classes_per_set)  # one hot encode

            [b, num_classes, spc, h, w, c] = self.support_set_images.get_shape().as_list()
            support_set_images_ = tf.reshape(self.support_set_images, shape=(b,  num_classes*spc, h, w, c))

            ## zero step: extractor feature embeddings
            target_image_ = tf.expand_dims(self.target_image, axis=1) #(b, 1, h, w, c)
            ## merge support set and target set, in order to share the feature extractors
            support_target_images = tf.concat([support_set_images_, target_image_], axis=1) #(b, n*k+1, h, w, c)
            print("+++ support_target images ", support_target_images.get_shape()) # (32, 6, 84, 84, 3)
            print("+++ support_target images [:-1]", support_target_images[:, :-1].get_shape()) # (32, 5, 84, 84, 3)
            support_target_embeddings = self.extractor(support_target_images, training=self.is_training, keep_prob=self.keep_prob)
            print("+++", support_target_embeddings.get_shape()) # (32, 6, 6, 6 , 96)  the last dimension is feature dimension

            ## first  step: generate task feature representation
            task_contexts = self.tce(support_target_embeddings[:, :-1], training=self.is_training) # (bs, num_task_features) (32, 64)

            ## second step: transform images via conditional meta task convolution
            ## todo In order to generate ensemble weights for the same task instance, we just need to run generation network several times
            ensemble_preds = []
            for m in range(M):
                trans_support_images_list = []
                trans_target_images_list = []
                for i, (tc, ste) in enumerate(zip(tf.unstack(task_contexts), tf.unstack(support_target_embeddings))):
                    print("============ In task instance ", i)
                    # support task image embeddings for one task
                    steb = self.Classifier(image_embedding=ste, task_context=tc, training=self.is_training, keep_prob=self.keep_prob)  #(6, 4608)
                    trans_support_images_list.append(steb[:-1])
                    trans_target_images_list.append(steb[-1])


                trans_support = tf.stack(trans_support_images_list)
                trans_target = tf.stack(trans_target_images_list)
                print("==" * 10)  # shape error
                print("trans support set shape and target shape ", trans_support.get_shape(), trans_target.get_shape())

                similarities = self.dn(support_set=trans_support, input_image=trans_target, name="distance_calculation",
                                       training=self.is_training)  # get similarity between support set embeddings and target

                preds = self.classify(similarities,
                                      support_set_y=self.support_set_labels_, name='classify', training=self.is_training)
                print("preds shape ", preds.get_shape())  # (bs, num_classes)
                ensemble_preds.append(tf.arg_max(preds, 1))

            ensemble_preds = tf.stack(ensemble_preds)


        return ensemble_preds 

def densenet_test():
    image_batch_placeholder = tf.placeholder(tf.float32, shape=[None, 224, 224, 3])
    label_batch_placeholder = tf.placeholder(tf.int64, shape=[BATCH_SIZE])
    if_training_placeholder = tf.placeholder(tf.bool, shape=[])

    image_batch, label_batch, filename_batch = data_provider.feed_data(if_random = False, if_training = False)
    label_batch_dense = tf.arg_max(label_batch, dimension = 1)

    if_training = tf.Variable(False, name='if_training', trainable=False)

    logits = tf.reshape(densenet.densenet_inference(image_batch_placeholder, if_training_placeholder, 1.0), [BATCH_SIZE, 5])
    logits_batch = tf.to_int64(tf.arg_max(logits, dimension = 1))

    correct_prediction = tf.equal(logits_batch, label_batch_placeholder)
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

    checkpoint = tf.train.get_checkpoint_state("./models")
    saver = tf.train.Saver()

    config = tf.ConfigProto()
    config.gpu_options.allow_growth=True
    with tf.Session(config=config) as sess:
        sess.run(tf.global_variables_initializer())
        tf.logging.info("Restoring full model from checkpoint file %s",checkpoint.model_checkpoint_path)
        saver.restore(sess, checkpoint.model_checkpoint_path)

        accuracy_accu = 0

        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord, sess = sess)

        for i in range(int(TEST_SET_SIZE / BATCH_SIZE)):
            image_out, label_batch_dense_out, filename_out = sess.run([image_batch, label_batch_dense, filename_batch])
            print("label: ", label_batch_dense_out)
            accuracy_out, infer_out = sess.run([accuracy, logits_batch], feed_dict={image_batch_placeholder: image_out,
                                                                                    label_batch_placeholder: label_batch_dense_out,
                                                                                    if_training_placeholder: if_training})
            accuracy_out = np.asarray(accuracy_out)
            print("infer: ", infer_out)
            print(' ')
            accuracy_accu = accuracy_out + accuracy_accu

        print(accuracy_accu / TEST_SET_SIZE * BATCH_SIZE)

        tf.train.write_graph(sess.graph_def, 'graph/', 'my_graph.pb', as_text=False)

        coord.request_stop()
        coord.join(threads)
        sess.close()
    return 0 

def add_loss_op(self, q, target_q):
        """
        Sets the loss of a batch, self.loss is a scalar

        Args:
            q: (tf tensor) shape = (batch_size, num_actions)
            target_q: (tf tensor) shape = (batch_size, num_actions)
        """
        # you may need this variable
        num_actions = self.env.action_space.n

        ##############################################################
        """
        TODO: The loss for an example is defined as:
                Q_samp(s) = r if done
                          = r + gamma * Q_target(s', max_a'Q(s',a'))
                loss = (Q_samp(s) - Q(s, a))^2 

              You need to compute the average of the loss over the minibatch
              and store the resulting scalar into self.loss

        HINT: - config variables are accessible through self.config
              - you can access placeholders like self.a (for actions)
                self.r (rewards) or self.done_mask for instance
              - you may find the following functions useful
                    - tf.cast
                    - tf.reduce_max / reduce_sum
                    - tf.one_hot
                    - ...

        (be sure that you set self.loss)
        """
        ##############################################################
        ##################### YOUR CODE HERE - 4-5 lines #############

        #done = tf.cast(self.done_mask, tf.float32)
        idx = tf.arg_max(q, dimension=1)
        idx_one_hot = tf.one_hot(idx, num_actions)
        temp = self.r + self.config.gamma*tf.reduce_sum(tf.multiply(target_q, idx_one_hot), axis=1)
        q_samp = tf.where(self.done_mask, self.r, temp)
        action = tf.one_hot(self.a, num_actions)
        q_new = tf.reduce_sum(tf.multiply(action,q), axis=1)
        self.loss = tf.reduce_mean(tf.square(q_new - q_samp))

        ##############################################################
        ######################## END YOUR CODE #######################