import tensorflow as tf

import tensorlayer as tl

import argparse

from data.mx2tfrecords import parse_function

import os

# from nets.L_Resnet_E_IR import get_resnet

# from nets.L_Resnet_E_IR_GBN import get_resnet

from nets.L_Resnet_E_IR_fix_issue9 import get_resnet

from losses.face_losses import arcface_loss

from tensorflow.core.protobuf import config_pb2

import time

from data.eval_data_reader import load_bin

from verification import ver_test

def get_parser():

return args

if __name__ == '__main__':

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# 1. define global parameters

args = get_parser()

global_step = tf.Variable(name='global_step', initial_value=0, trainable=False)

inc_op = tf.assign_add(global_step, 1, name='increment_global_step')

images = tf.placeholder(name='img_inputs', shape=[None, *args.image_size, 3], dtype=tf.float32)

labels = tf.placeholder(name='img_labels', shape=[None, ], dtype=tf.int64)

# trainable = tf.placeholder(name='trainable_bn', dtype=tf.bool)

dropout_rate = tf.placeholder(name='dropout_rate', dtype=tf.float32)

# 2 prepare train datasets and test datasets by using tensorflow dataset api

# 2.1 train datasets

# the image is substracted 127.5 and multiplied 1/128.

# random flip left right

tfrecords_f = os.path.join(args.tfrecords_file_path, 'tran.tfrecords')

dataset = tf.data.TFRecordDataset(tfrecords_f)

dataset = dataset.map(parse_function)

dataset = dataset.shuffle(buffer_size=args.buffer_size)

dataset = dataset.batch(args.batch_size)

iterator = dataset.make_initializable_iterator()

next_element = iterator.get_next()

# 2.2 prepare validate datasets

ver_list = []

ver_name_list = []

for db in args.eval_datasets:

print('begin db %s convert.' % db)

data_set = load_bin(db, args.image_size, args)

ver_list.append(data_set)

ver_name_list.append(db)

# 3. define network, loss, optimize method, learning rate schedule, summary writer, saver

# 3.1 inference phase

w_init_method = tf.contrib.layers.xavier_initializer(uniform=False)

net = get_resnet(images, args.net_depth, type='ir', w_init=w_init_method, trainable=True, keep_rate=dropout_rate)

# 3.2 get arcface loss

logit = arcface_loss(embedding=net.outputs, labels=labels, w_init=w_init_method, out_num=args.num_output)

# test net because of batch normal layer

tl.layers.set_name_reuse(True)

test_net = get_resnet(images, args.net_depth, type='ir', w_init=w_init_method, trainable=False, reuse=True, keep_rate=dropout_rate)

embedding_tensor = test_net.outputs

# 3.3 define the cross entropy

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

# inference_loss_avg = tf.reduce_mean(inference_loss)

# 3.4 define weight deacy losses

# for var in tf.trainable_variables():

# print(var.name)

# print('##########'*30)

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_deacy)(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_deacy)(W)

for weights in tl.layers.get_variables_with_name('embedding_weights', True, True):

wd_loss += tf.contrib.layers.l2_regularizer(args.weight_deacy)(weights)

for gamma in tl.layers.get_variables_with_name('gamma', True, True):

wd_loss += tf.contrib.layers.l2_regularizer(args.weight_deacy)(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)

# for bias in tl.layers.get_variables_with_name('resnet_v1_50/E_DenseLayer/b', True, True):

# wd_loss += tf.contrib.layers.l2_regularizer(args.weight_deacy)(bias)

# 3.5 total losses

total_loss = inference_loss + wd_loss

# 3.6 define the learning rate schedule

p = int(512.0/args.batch_size)

lr_steps = [p*val for val in args.lr_steps]

print(lr_steps)

lr = tf.train.piecewise_constant(global_step, boundaries=lr_steps, values=[0.001, 0.0005, 0.0003, 0.0001], name='lr_schedule')

# 3.7 define the optimize method

opt = tf.train.MomentumOptimizer(learning_rate=lr, momentum=args.momentum)

# 3.8 get train op

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)

# train_op = opt.minimize(total_loss, global_step=global_step)

# 3.9 define the inference accuracy used during validate or test

pred = tf.nn.softmax(logit)

acc = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(pred, axis=1), labels), dtype=tf.float32))

# 3.10 define sess

config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=args.log_device_mapping)

config.gpu_options.allow_growth = True

sess = tf.Session(config=config)

# 3.11 summary writer

summary = tf.summary.FileWriter(args.summary_path, sess.graph)

summaries = []

# # 3.11.1 add grad histogram op

for grad, var in grads:

if grad is not None:

summaries.append(tf.summary.histogram(var.op.name + '/gradients', grad))

# 3.11.2 add trainabel variable gradients

for var in tf.trainable_variables():

summaries.append(tf.summary.histogram(var.op.name, var))

# 3.11.3 add loss summary

summaries.append(tf.summary.scalar('inference_loss', inference_loss))

summaries.append(tf.summary.scalar('wd_loss', wd_loss))

summaries.append(tf.summary.scalar('total_loss', total_loss))

# 3.11.4 add learning rate

summaries.append(tf.summary.scalar('leraning_rate', lr))

summary_op = tf.summary.merge(summaries)

# 3.12 saver

saver = tf.train.Saver(max_to_keep=args.saver_maxkeep)

# 3.13 init all variables

sess.run(tf.global_variables_initializer())

# restore_saver = tf.train.Saver()

# restore_saver.restore(sess, '/home/aurora/workspaces2018/InsightFace_TF/output/ckpt/InsightFace_iter_1110000.ckpt')

# 4 begin iteration

if not os.path.exists(args.log_file_path):

os.makedirs(args.log_file_path)

log_file_path = args.log_file_path + '/train' + time.strftime('_%Y-%m-%d-%H-%M', time.localtime(time.time())) + '.log'

log_file = open(log_file_path, 'w')

# 4 begin iteration

count = 0

total_accuracy = {}

for i in range(args.epoch):

sess.run(iterator.initializer)

while True:

try:

images_train, labels_train = sess.run(next_element)

feed_dict = {images: images_train, labels: labels_train, dropout_rate: 0.4}

feed_dict.update(net.all_drop)

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=feed_dict,

options=config_pb2.RunOptions(report_tensor_allocations_upon_oom=True))

end = time.time()

pre_sec = args.batch_size/(end - start)

# print training information

if count > 0 and count % args.show_info_interval == 0:

print('epoch %d, total_step %d, total loss is %.2f , inference loss is %.2f, weight deacy '

'loss is %.2f, training accuracy is %.6f, time %.3f samples/sec' %

(i, count, total_loss_val, inference_loss_val, wd_loss_val, acc_val, pre_sec))

count += 1

# save summary

if count > 0 and count % args.summary_interval == 0:

feed_dict = {images: images_train, labels: labels_train, dropout_rate: 0.4}

feed_dict.update(net.all_drop)

summary_op_val = sess.run(summary_op, feed_dict=feed_dict)

summary.add_summary(summary_op_val, count)

# 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)

# validate

if count > 0 and count % args.validate_interval == 0:

feed_dict_test ={dropout_rate: 1.0}

feed_dict_test.update(tl.utils.dict_to_one(net.all_drop))

results = ver_test(ver_list=ver_list, ver_name_list=ver_name_list, nbatch=count, sess=sess,

embedding_tensor=embedding_tensor, batch_size=args.batch_size, feed_dict=feed_dict_test,

input_placeholder=images)

print('test accuracy is: ', str(results[0]))

total_accuracy[str(count)] = results[0]

log_file.write('########'*10+'\n')

log_file.write(','.join(list(total_accuracy.keys())) + '\n')

log_file.write(','.join([str(val) for val in list(total_accuracy.values())])+'\n')

log_file.flush()

if max(results) > 0.996:

print('best accuracy is %.5f' % max(results))

filename = 'InsightFace_iter_best_{:d}'.format(count) + '.ckpt'

filename = os.path.join(args.ckpt_path, filename)

saver.save(sess, filename)

log_file.write('######Best Accuracy######'+'\n')

log_file.write(str(max(results))+'\n')

log_file.write(filename+'\n')

log_file.flush()

except tf.errors.OutOfRangeError:

print("End of epoch %d" % i)

break

log_file.close()

log_file.write('\n')