if __name__ == '__main__':
  os.environ['CUDA_VISIBLE_DEVICES'] = "0"
  args = get_parser()
  train_set, train_label = msm.get_train_set(args.dataset_dir, args.clean_list_path)
  test_set_left, test_set_right, test_set_label = fe.get_test_set()

  global_step = tf.Variable(name='global_step', initial_value=0, trainable=False)
  inc_op = tf.assign_add(global_step, 1, name='increment_global_step')
  x = tf.placeholder(tf.float32, [None, 112, 112, 3], name='x')
  y = tf.placeholder(tf.int64, [None, ], name='label')
  dropout_rate = tf.placeholder(name='dropout_rate', dtype=tf.float32)

  w_init_method = tf.contrib.layers.xavier_initializer(uniform=False)
  network = resnet.get_resnet(x, 50, type='ir', w_init=w_init_method, trainable=True, keep_rate=dropout_rate)
  logit   = resnet.arcface_loss(embedding=network.outputs, labels=y, w_init=w_init_method, out_num=args.num_output)

  tl.layers.set_name_reuse(True)
  test_net = resnet.get_resnet(x, 50, type='ir', w_init=w_init_method, trainable=False, reuse=True, keep_rate=dropout_rate)
  embedding_tensor = test_net.outputs

  inference_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logit, labels=y))

  wd_loss = 0
  for weights in tl.layers.get_variables_with_name('W_conv2d', True, True):
    wd_loss += tf.contrib.layers.l2_regularizer(args.weight_decay)(weights)
  for W in tl.layers.get_variables_with_name('resnet_v1_50/E_DenseLayer/W', True, True):
    wd_loss += tf.contrib.layers.l2_regularizer(args.weight_decay)(W)
  for weights in tl.layers.get_variables_with_name('embedding_weights', True, True):
    wd_loss += tf.contrib.layers.l2_regularizer(args.weight_decay)(weights)
  for gamma in tl.layers.get_variables_with_name('gamma', True, True):
    wd_loss += tf.contrib.layers.l2_regularizer(args.weight_decay)(gamma)
  # for beta in tl.layers.get_variables_with_name('beta', True, True):
  #     wd_loss += tf.contrib.layers.l2_regularizer(args.weight_deacy)(beta)
  for alphas in tl.layers.get_variables_with_name('alphas', True, True):
    wd_loss += tf.contrib.layers.l2_regularizer(args.weight_deacy)(alphas)

  total_loss = inference_loss + wd_loss

  p = int(512.0/args.batch_size)
  lr_steps = [p*val for val in args.lr_steps]
  lr = tf.train.piecewise_constant(global_step, boundaries=lr_steps, values=[0.001, 0.0005, 0.0003, 0.0001], name='lr_schedule')

  opt = tf.train.MomentumOptimizer(learning_rate=lr, momentum=args.momentum)
  grads = opt.compute_gradients(total_loss)
  update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
  with tf.control_dependencies(update_ops):
    train_op = opt.apply_gradients(grads, global_step=global_step)

  pred = tf.nn.softmax(logit)
  acc = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(pred, axis=1), y), dtype=tf.float32))

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

  with tf.Session(config=config) as sess:
    saver = tf.train.Saver(max_to_keep=args.saver_maxkeep)
    sess.run(tf.global_variables_initializer())

    idx = 0
    count = 0
    for i in range(args.epoch):
      batch_train, batch_label, idx = get_train_batch(idx, args.batch_size, train_set, train_label)

      start = time.time()
      _, total_loss_val, inference_loss_val, wd_loss_val, _, acc_val = sess.run([train_op, total_loss, inference_loss, wd_loss, inc_op, acc], feed_dict={x: batch_train, y:batch_label, dropout_rate:0.4})
      end   = time.time()
      pre_sec = args.batch_size / (end - start)

      if count > 0 and count % args.show_info_interval == 0:
        print('epoch %d, total_step %d, total loss %.2f, inference loss %.2f, training acc %.6f, time %.3f' % (i, count, total_loss_val, inference_loss_val, acc_val, pre_sec))
        count += 1

                      # save ckpt files
      if count > 0 and count % args.ckpt_interval == 0:
        filename = 'InsightFace_iter_{:d}'.format(count) + '.ckpt'
        filename = os.path.join(args.ckpt_path, filename)
        saver.save(sess, filename)