text: import os
  import math
  
  import numpy as np
  import pandas as pd
  import tensorflow as tf
  from keras.utils import to_categorical
  from sklearn.metrics import recall_score as recall
  from sklearn.metrics import confusion_matrix as confusion
  
  
  def attention(inputs, attention_size, time_major=False, return_alphas=False):
      """
      Adapted from
          https://github.com/xuanjihe/speech-emotion-recognition/blob/master/attention.py
      Note that the author has changed the activation function of the fully
      connected layer from tanh to sigmoid.
  
      Attention mechanism layer which reduces RNN/Bi-RNN outputs with Attention
      vector.
  
      The idea was proposed in the article by Z. Yang et al., "Hierarchical
      Attention Networks for Document Classification",
      2016: http://www.aclweb.org/anthology/N16-1174.
      Variables notation is also inherited from the article
  
      Args:
          inputs: The Attention inputs.
              Matches outputs of RNN/Bi-RNN layer (not final state):
                  In case of RNN, this must be RNN outputs `Tensor`:
                      If time_major == False (default), this must be a tensor of
                      shape:
                          `[batch_size, max_time, cell.output_size]`.
                      If time_major == True, this must be a tensor of shape:
                          `[max_time, batch_size, cell.output_size]`.
                  In case of Bidirectional RNN, this must be a tuple
                  (outputs_fw, outputs_bw) containing the forward and the
                  backward RNN outputs `Tensor`.
                      If time_major == False (default),
                          outputs_fw is a `Tensor` shaped:
                          `[batch_size, max_time, cell_fw.output_size]`
                          and outputs_bw is a `Tensor` shaped:
                          `[batch_size, max_time, cell_bw.output_size]`.
                      If time_major == True,
                          outputs_fw is a `Tensor` shaped:
                          `[max_time, batch_size, cell_fw.output_size]`
                          and outputs_bw is a `Tensor` shaped:
                          `[max_time, batch_size, cell_bw.output_size]`.
          attention_size: Linear size of the Attention weights.
          time_major: The shape format of the `inputs` Tensors.
              If true, these `Tensors` must be shaped
              `[max_time, batch_size, depth]`.
              If false, these `Tensors` must be shaped
              `[batch_size, max_time, depth]`.
              Using `time_major = True` is a bit more efficient because it avoids
              transposes at the beginning and end of the RNN calculation.
              However, most TensorFlow data is batch-major, so by default this
              function accepts input and emits output in batch-major form.
          return_alphas: Whether to return attention coefficients variable along
              with layer's output. Used for visualization purpose.
      Returns:
          The Attention output `Tensor`.
          In case of RNN, this will be a `Tensor` shaped:
              `[batch_size, cell.output_size]`.
          In case of Bidirectional RNN, this will be a `Tensor` shaped:
              `[batch_size, cell_fw.output_size + cell_bw.output_size]`.
      """
  
      if isinstance(inputs, tuple):
          # In case of Bi-RNN, concatenate the forward and backward RNN outputs.
          inputs = tf.concat(inputs, 2)
  
      if time_major:
          # (T,B,D) => (B,T,D)
          inputs = tf.array_ops.transpose(inputs, [1, 0, 2])
  
      hidden_size = inputs.shape[2].value  # D value - hidden size of RNN layer
  
      # Trainable parameters
      with tf.variable_scope("attention"):
          W_omega = tf.Variable(
              tf.random_normal([hidden_size, attention_size], stddev=0.1),
              name="weight")
          b_omega = tf.Variable(
              tf.random_normal([attention_size], stddev=0.1), name="bias")
          u_omega = tf.Variable(
              tf.random_normal([attention_size], stddev=0.1), name="u")
  
      # Applying fully connected layer with non-linear activation to each of
      # the B*T timestamps; the shape of `v` is (B,T,D)*(D,A)=(B,T,A),
      # where A=attention_size
      v = tf.sigmoid(tf.tensordot(inputs, W_omega, axes=1) + b_omega)
      # For each of the timestamps its vector of size A from `v` is reduced with
      # `u` vector
      vu = tf.tensordot(v, u_omega, axes=1)   # (B,T) shape
      alphas = tf.nn.softmax(vu)              # (B,T) shape also
  
      # Output of (Bi-)RNN is reduced with attention vector; the result has (B,D)
      # shape
      output = tf.reduce_sum(inputs * tf.expand_dims(alphas, -1), 1)
  
      if not return_alphas:
          return output
      else:
          return output, alphas
  
  
  def leaky_relu(x, leakiness=0.0):
      return tf.where(tf.less(x, 0.0), leakiness * x, x, name='leaky_relu')
  
  
  def batch_norm_wrapper(inputs, is_training, decay=0.999, epsilon=1e-3):
      with tf.variable_scope("batch_norm"):
          scale = tf.Variable(tf.ones([inputs.get_shape()[-1]]), name="scale")
          beta = tf.Variable(tf.zeros([inputs.get_shape()[-1]]), name="beta")
          pop_mean = tf.Variable(tf.zeros([inputs.get_shape()[-1]]),
                                 name="mean", trainable=False)
          pop_var = tf.Variable(tf.ones([inputs.get_shape()[-1]]),
                                name="var", trainable=False)
  
      if is_training is not None:
          batch_mean, batch_var = tf.nn.moments(inputs, [0])
          train_mean = tf.assign(
              pop_mean, pop_mean * decay + batch_mean * (1 - decay))
          train_var = tf.assign(
              pop_var, pop_var * decay + batch_var * (1 - decay))
          with tf.control_dependencies([train_mean, train_var]):
              return tf.nn.batch_normalization(
                  inputs, batch_mean, batch_var, beta, scale, epsilon)
      else:
          return tf.nn.batch_normalization(
              inputs, pop_mean, pop_var, beta, scale, epsilon)
  
  
  def acrnn(inputs, num_classes=4, is_training=True, dropout_keep_prob=1,
            num_filt_1=128, num_filt_2=256, mp_fsize=(2, 4), num_lstm_units=128,
            num_linear_units=768, num_fcn_units=64):
      """
      Reference: "3-D Convolutional Recurrent Neural Networks with Attention
      Model for Speech Emotion Recognition"
      Authors: Chen Mingyi and
               He Xuanji and
               Yang Jing and
               Zhang Han.
  
      Adapted from
          https://github.com/xuanjihe/speech-emotion-recognition/blob/master/acrnn1.py
  
      3D-CRNN model wrapper with Tensorflow. Take `inputs` as model input and
      return the model logits.
  
      Parameters
      ----------
      inputs : tensor
          Input tensor.
      num_classes : int
          Number of classes.
      is_training : bool
          Whether the model is built for training or not.
      dropout_keep_prob : float
          Probability of keeping connections in dropout.
  
      num_filt_1 : int
          Number of filters of the first convolutional layer.
      num_filt_2 : int
          Number of filters of all the rest convolutional layers.
      mp_fsize : tuple or list
          Max pooling filter size (tuple/list of height and width, respectively)
  
      num_lstm_units : int
          Number of neurons in LSTM cell.
      num_linear_units : int
          Number of units in the linear layer (the layer after the last
          convolutional layer).
      num_fcn_units : int
          Number of neurons in the last fully connected layer.
  
      Returns
      -------
      logits
          Logits of the model.
  
      """
      image_height, image_width = inputs.get_shape().as_list()[1:3]
      mp_fsize_h, mp_fsize_w = mp_fsize
      # `time_step` and `p` can be calculated when knowing maxpool filter size
      time_step = image_height // mp_fsize_h
      p = image_width // mp_fsize_w
  
      layer1_filter = tf.get_variable(
          'layer1_filter', shape=[5, 3, 3, num_filt_1], dtype=tf.float32,
          initializer=tf.truncated_normal_initializer(stddev=0.1))
      layer1_bias = tf.get_variable(
          'layer1_bias', shape=[num_filt_1], dtype=tf.float32,
          initializer=tf.constant_initializer(0.1))
      layer1_stride = [1, 1, 1, 1]
  
      layer2_filter = tf.get_variable(
          'layer2_filter', dtype=tf.float32,
          shape=[5, 3, num_filt_1, num_filt_2],
          initializer=tf.truncated_normal_initializer(stddev=0.1))
      layer2_bias = tf.get_variable(
          'layer2_bias', shape=[num_filt_2], dtype=tf.float32,
          initializer=tf.constant_initializer(0.1))
      layer2_stride = [1, 1, 1, 1]
  
      layer3_filter = tf.get_variable(
          'layer3_filter', dtype=tf.float32,
          shape=[5, 3, num_filt_2, num_filt_2],
          initializer=tf.truncated_normal_initializer(stddev=0.1))
      layer3_bias = tf.get_variable(
          'layer3_bias', shape=[num_filt_2], dtype=tf.float32,
          initializer=tf.constant_initializer(0.1))
      layer3_stride = [1, 1, 1, 1]
  
      layer4_filter = tf.get_variable(
          'layer4_filter', dtype=tf.float32,
          shape=[5, 3, num_filt_2, num_filt_2],
          initializer=tf.truncated_normal_initializer(stddev=0.1))
      layer4_bias = tf.get_variable(
          'layer4_bias', shape=[num_filt_2], dtype=tf.float32,
          initializer=tf.constant_initializer(0.1))
      layer4_stride = [1, 1, 1, 1]
  
      layer5_filter = tf.get_variable(
          'layer5_filter', dtype=tf.float32,
          shape=[5, 3, num_filt_2, num_filt_2],
          initializer=tf.truncated_normal_initializer(stddev=0.1))
      layer5_bias = tf.get_variable(
          'layer5_bias', shape=[num_filt_2], dtype=tf.float32,
          initializer=tf.constant_initializer(0.1))
      layer5_stride = [1, 1, 1, 1]
  
      layer6_filter = tf.get_variable(
          'layer6_filter', dtype=tf.float32,
          shape=[5, 3, num_filt_2, num_filt_2],
          initializer=tf.truncated_normal_initializer(stddev=0.1))
      layer6_bias = tf.get_variable(
          'layer6_bias', shape=[num_filt_2], dtype=tf.float32,
          initializer=tf.constant_initializer(0.1))
      layer6_stride = [1, 1, 1, 1]
  
      linear1_weight = tf.get_variable(
          'linear1_weight', shape=[p * num_filt_2, num_linear_units],
          dtype=tf.float32, initializer=tf.truncated_normal_initializer(
              stddev=0.1))
      linear1_bias = tf.get_variable(
          'linear1_bias', shape=[num_linear_units], dtype=tf.float32,
          initializer=tf.constant_initializer(0.1))
  
      fully1_weight = tf.get_variable(
          'fully1_weight', shape=[2 * num_lstm_units, num_fcn_units],
          dtype=tf.float32, initializer=tf.truncated_normal_initializer(
              stddev=0.1))
      fully1_bias = tf.get_variable(
          'fully1_bias', shape=[num_fcn_units], dtype=tf.float32,
          initializer=tf.constant_initializer(0.1))
      fully2_weight = tf.get_variable(
          'fully2_weight', shape=[num_fcn_units, num_classes], dtype=tf.float32,
          initializer=tf.truncated_normal_initializer(stddev=0.1))
      fully2_bias = tf.get_variable(
          'fully2_bias', shape=[num_classes], dtype=tf.float32,
          initializer=tf.constant_initializer(0.1))
  
      layer1 = tf.nn.conv2d(inputs, layer1_filter, layer1_stride, padding='SAME')
      layer1 = tf.nn.bias_add(layer1, layer1_bias)
      layer1 = leaky_relu(layer1, 0.01)
      layer1 = tf.nn.max_pool(
          layer1, ksize=[1, mp_fsize_h, mp_fsize_w, 1],
          strides=[1, mp_fsize_h, mp_fsize_w, 1],
          padding='VALID', name='max_pool')
      layer1 = tf.contrib.layers.dropout(
          layer1, keep_prob=dropout_keep_prob, is_training=is_training)
  
      layer2 = tf.nn.conv2d(layer1, layer2_filter, layer2_stride, padding='SAME')
      layer2 = tf.nn.bias_add(layer2, layer2_bias)
      layer2 = leaky_relu(layer2, 0.01)
      layer2 = tf.contrib.layers.dropout(
          layer2, keep_prob=dropout_keep_prob, is_training=is_training)
  
      layer3 = tf.nn.conv2d(layer2, layer3_filter, layer3_stride, padding='SAME')
      layer3 = tf.nn.bias_add(layer3, layer3_bias)
      layer3 = leaky_relu(layer3, 0.01)
      layer3 = tf.contrib.layers.dropout(
          layer3, keep_prob=dropout_keep_prob, is_training=is_training)
  
      layer4 = tf.nn.conv2d(layer3, layer4_filter, layer4_stride, padding='SAME')
      layer4 = tf.nn.bias_add(layer4, layer4_bias)
      layer4 = leaky_relu(layer4, 0.01)
      layer4 = tf.contrib.layers.dropout(
          layer4, keep_prob=dropout_keep_prob, is_training=is_training)
  
      layer5 = tf.nn.conv2d(layer4, layer5_filter, layer5_stride, padding='SAME')
      layer5 = tf.nn.bias_add(layer5, layer5_bias)
      layer5 = leaky_relu(layer5, 0.01)
      layer5 = tf.contrib.layers.dropout(
          layer5, keep_prob=dropout_keep_prob, is_training=is_training)
  
      layer6 = tf.nn.conv2d(layer5, layer6_filter, layer6_stride, padding='SAME')
      layer6 = tf.nn.bias_add(layer6, layer6_bias)
      layer6 = leaky_relu(layer6, 0.01)
      layer6 = tf.contrib.layers.dropout(
          layer6, keep_prob=dropout_keep_prob, is_training=is_training)
  
      layer6 = tf.reshape(layer6, [-1, p * num_filt_2])
  
      linear1 = tf.matmul(layer6, linear1_weight) + linear1_bias
      linear1 = batch_norm_wrapper(linear1, is_training)
      linear1 = leaky_relu(linear1, 0.01)
      linear1 = tf.reshape(linear1, [-1, time_step, num_linear_units])
  
      # Define lstm cells with tensorflow
      # Forward direction cell
      gru_fw_cell = tf.contrib.rnn.BasicLSTMCell(num_lstm_units, forget_bias=1.0)
      # Backward direction cell
      gru_bw_cell = tf.contrib.rnn.BasicLSTMCell(num_lstm_units, forget_bias=1.0)
  
      # Now we feed `layer_3` into the LSTM BRNN cell and obtain the LSTM BRNN
      # output.
      outputs, output_states = tf.nn.bidirectional_dynamic_rnn(
          cell_fw=gru_fw_cell, cell_bw=gru_bw_cell, inputs=linear1,
          dtype=tf.float32, time_major=False, scope='LSTM')
  
      # Attention layer
      gru, alphas = attention(outputs, 1, return_alphas=True)
      # Fully connected layers
      fully1 = tf.matmul(gru, fully1_weight) + fully1_bias
      fully1 = leaky_relu(fully1, 0.01)
      fully1 = tf.nn.dropout(fully1, dropout_keep_prob)
  
      logits = tf.add(
          tf.matmul(fully1, fully2_weight), fully2_bias, name="logits")
      return logits
  
  
  def CB_loss(labels, logits, samples_per_cls, beta=0.9999, is_training=True,
              name=None):
      """
      Reference: "Class-Balanced Loss Based on Effective Number of Samples"
      Authors: Yin Cui and
               Menglin Jia and
               Tsung Yi Lin and
               Yang Song and
               Serge J. Belongie
      https://arxiv.org/abs/1901.05555, CVPR'19.
  
      Adapted from the PyTorch version:
        https://github.com/vandit15/Class-balanced-loss-pytorch/blob/master/class_balanced_loss.py
  
      Compute the Class Balanced Loss between `logits` and the ground truth
      `labels`. Class Balanced Loss:
          ((1 - beta) / (1 - beta^n)) * Loss(labels, logits)
      where Loss is standard cross entropy loss.
  
      Parameters
      ----------
      labels : tensor
          A tensor of shape (batch_size, num_classes).
      logits : tensor
          A tensor of shape (batch_size, num_classes). Output probabilities of
          the model.
      samples_per_cls : list or ndarray
          A list or ndarray of length `num_classes`.
      beta : float
          Hyperparameter for Class Balanced Loss.
      is_training : bool
  
      Returns
      -------
      loss : tensor
          If training, return scalar reduced loss over non-zero weights.
          Otherwise, return un-reduced loss (whose shape is (batch_size,)).
  
      """
      num_classes = labels.get_shape().as_list()[1]
  
      effective_num = 1.0 - np.power(beta, samples_per_cls)
      weights = (1.0 - beta) / np.array(effective_num)
      weights = weights / np.sum(weights) * num_classes
      weights = np.expand_dims(weights, 0)
  
      weights = tf.multiply(weights, labels)
      weights = tf.math.reduce_sum(weights, axis=1)
      weights = tf.expand_dims(weights, axis=-1)
  
      softmax = tf.nn.softmax(logits, axis=-1)
  
      if is_training:
          loss = tf.losses.log_loss(
              labels=labels, predictions=softmax, weights=weights)
          loss = tf.identity(loss, name=name)
      else:
          loss = tf.losses.log_loss(
              labels=labels, predictions=softmax, weights=weights,
              reduction=tf.losses.Reduction.NONE)
          loss = tf.reduce_sum(loss, axis=-1, name=name)
  
      return loss
  
  
  def test(test_data, test_labels, test_segs_labels, test_segs, sess,
           num_classes=4, batch_size=64):
      """Test a 3DCRNN model.
  
      Parameters
      ----------
      test_data: ndarray
      test_labels: ndarray
      test_segs_labels: ndarray
      test_segs: ndarray
          The first four arguments are the data in order obtained using the
          script `extract_mel.py`.
      num_classes : int
          Number of classes.
      batch_size : int
  
      Returns
      -------
      test_cost : float
          The cost value of the model on the test data.
      test_ua : float
          The unweighted accuracy of the model on the test data.
      test_conf : float
          The confusion matrix of the model on the test data.
      """
      # Parameters
      num_test = test_segs.shape[0]
      num_test_segs = test_data.shape[0]
      # Store predictions
      y_test_segs = np.empty(
          (num_test_segs, num_classes), dtype=np.float32)
      y_test = np.empty((num_test, num_classes), dtype=np.float32)
  
      test_cost = 0
      if num_test_segs < batch_size:
          loss, y_test_segs = sess.run(
              ["cross_entropy:0", "logits:0"],
              feed_dict={"X_input:0": test_data, "Y_input:0": test_segs_labels,
                         "is_training:0": False, "keep_prob:0": 1}
          )
          test_cost = test_cost + np.sum(loss)
  
      num_batches_test = divmod((num_test_segs), batch_size)[0]
      for batch_test in range(num_batches_test):
          v_begin = batch_test * batch_size
          v_end = (batch_test + 1) * batch_size
          if batch_test == num_batches_test - 1:
              if v_end < num_test_segs:
                  v_end = num_test_segs
          loss, y_test_segs[v_begin:v_end] = sess.run(
              ["cross_entropy:0", "logits:0"],
              feed_dict={"X_input:0": test_data[v_begin:v_end],
                         "Y_input:0": test_segs_labels[v_begin:v_end],
                         "is_training:0": False, "keep_prob:0": 1})
          test_cost = test_cost + np.sum(loss)
      test_cost = test_cost / num_test_segs
      # Accumulate results, since each utterance might contain
      # more than one segments
      curr_i = 0
      for v in range(num_test):
          y_test[v, :] = np.max(
              y_test_segs[curr_i:curr_i + test_segs[v]], 0)
          curr_i = curr_i + test_segs[v]
      # Accuracy and confusion matrix
      test_labels = np.argmax(test_labels, 1)
      y_test = np.argmax(y_test, 1)
      test_ua = (test_labels == y_test).mean()
      test_conf = confusion(test_labels, y_test)
      # Make confusion matrix into data frame for readability
      test_conf = pd.DataFrame({"ang": test_conf[:, 0], "sad": test_conf[:, 1],
                                "hap": test_conf[:, 2], "neu": test_conf[:, 3]})
      test_conf = test_conf.to_string(index=False)
      return test_cost, test_ua, test_conf
  
  
  def train(data, steps, batch_size, learning_rate, validate_every=10,
            random_seed=136, num_classes=4, grad_clip=False, dropout_keep_prob=1,
            save_path=None, use_CBL=False, beta=0.9999, perform_test=False,
            **kwargs):
      """Train a 3DCRNN model in an iterative-based manner with Tensorflow.
  
      Parameters
      ----------
      data : tuple
          Data extracted using `speech_utils.ACRNN.data_utils.extract_mel`.
      steps : int
          Number of global steps.
      batch_size : int
      learning_rate : float
      validate_every : int
          Number of training steps between each validation.
      random_seed : int
          Random seed for reproducibility.
      num_classes : int
          Number of classes.
      grad_clip : boolean
          Whether to clip gradients of Adam optimizer.
      dropout_keep_prob : float
          The probability of keeping a connection in dropout.
      save_path : str
          Path to save the best models.
      use_CBL : boolean
          Whether to use Class Balanced Loss.
      beta : float
          Hyperparameter for Class Balanced Loss. Default: 0.9999
      perform_test: bool
          Whether to test on test data at the end of training process.
      **kwargs
          Custom keyword arguments to pass to the ACRNN constructor.
  
      """
      # For reproducibility
      tf.random.set_random_seed(random_seed)
      np.random.seed(random_seed)
      # Load data
      train_data, train_labels = data[0:2]
      val_data, val_labels, val_segs_labels, val_segs = data[2:6]
      test_data, test_labels, test_segs_labels, test_segs = data[6:10]
      # One-hot encoding
      train_labels = to_categorical(train_labels, num_classes)
      val_labels = to_categorical(val_labels, num_classes)
      val_segs_labels = to_categorical(val_segs_labels, num_classes)
      test_labels = to_categorical(test_labels, num_classes)
      test_segs_labels = to_categorical(test_segs_labels, num_classes)
      # Parameters
      num_train, image_height, image_width, image_channel = train_data.shape
      best_val_ua = 0
      # Construct model
      X = tf.placeholder(
          tf.float32, name="X_input",
          shape=[None, image_height, image_width, image_channel])
      Y = tf.placeholder(tf.int32, name="Y_input", shape=[None, num_classes])
      is_training = tf.placeholder(tf.bool, name="is_training")
      lr = tf.placeholder(tf.float32)
      keep_prob = tf.placeholder(tf.float32, name="keep_prob")
      # Cost
      logits = acrnn(
          X, num_classes=num_classes, is_training=is_training,
          dropout_keep_prob=keep_prob, **kwargs)
      if not use_CBL:
          cross_entropy = tf.nn.softmax_cross_entropy_with_logits(
              labels=Y, logits=logits, name="cross_entropy")
          cost = tf.reduce_mean(cross_entropy)
      else:
          _, samples_per_cls = np.unique(
              np.argmax(train_labels, axis=-1), return_counts=True)
          cross_entropy = CB_loss(
              labels=tf.cast(Y, "float64"), logits=logits, is_training=False,
              samples_per_cls=samples_per_cls, beta=beta, name="cross_entropy")
          cost = CB_loss(labels=tf.cast(Y, "float64"), logits=logits,
                         samples_per_cls=samples_per_cls, beta=beta)
  
      var_trainable_op = tf.trainable_variables()
      # Optimizer
      if not grad_clip:
          # Do not apply gradient clipping
          train_op = tf.train.AdamOptimizer(lr).minimize(cost)
      else:
          # Apply gradient clipping
          grads, _ = tf.clip_by_global_norm(
              tf.gradients(cost, var_trainable_op), 5)
          opti = tf.train.AdamOptimizer(lr)
          train_op = opti.apply_gradients(zip(grads, var_trainable_op))
      # Predictions and accuracy
      correct_pred = tf.equal(tf.argmax(logits, 1), tf.argmax(Y, 1))
      accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32), name="acc")
      saver = tf.train.Saver(tf.global_variables())
      init = tf.global_variables_initializer()
      # Start training
      with tf.Session() as sess:
          sess.run(init)
          for i in range(steps):
              start = (i * batch_size) % num_train
              end = min(start + batch_size, num_train)
              feed_dict = {X: train_data[start:end], Y: train_labels[start:end],
                           is_training: True, keep_prob: dropout_keep_prob,
                           lr: learning_rate}
              _, train_cost, train_acc = sess.run([train_op, cost, accuracy],
                                                  feed_dict=feed_dict)
              # Test on validation set
              if i % validate_every == 0:
                  val_cost, val_ua, val_conf = test(
                      val_data, val_labels, val_segs_labels, val_segs, sess,
                      num_classes=num_classes, batch_size=batch_size)
                  # Update
                  if val_ua > best_val_ua:
                      best_val_ua = val_ua
                      best_val_conf = val_conf
                      best_val_step = i + 1
                      if save_path is not None:
                          saver.save(sess, save_path, global_step=i + 1)
  
                  print("Epoch: {}".format(i + 1))
                  print("Training cost: {:.04f}".format(train_cost))
                  print("Training UA: {:.02f}%".format(train_acc * 100))
                  print()
                  print("Valid cost: {:.04f}".format(val_cost))
                  print("Valid UA: {:.02f}%".format(val_ua * 100))
                  print('Valid confusion matrix:')
                  print(val_conf)
                  print()
                  print("Best valid UA: {:.02f}%".format(best_val_ua * 100))
                  print('Best valid confusion matrix:')
                  print(best_val_conf)
                  print("*" * 30)
  
          # Test on test set
          if perform_test:
              save_dir, _ = os.path.split(save_path)
              # Load weights
              saver.restore(sess, tf.train.latest_checkpoint(save_dir))
              # Start testing
              test_cost, test_ua, test_conf = test(
                  test_data, test_labels, test_segs_labels, test_segs, sess,
                  num_classes=num_classes, batch_size=batch_size)
              print("*" * 30)
              print("RESULTS ON TEST SET:")
              print("Test cost: {:.04f}, test unweighted accuracy: "
                    "{:.2f}%".format(test_cost, test_ua * 100))
              print('Test confusion matrix:')
              print(test_conf)
  