def train(): is_training = True # data pipeline imgs, true_boxes = gen_data_batch(cfg.data_path, cfg.batch_size*cfg.train.num_gpus) imgs_split = tf.split(imgs, cfg.train.num_gpus) true_boxes_split = tf.split(true_boxes, cfg.train.num_gpus) global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0.), trainable=False) lr = tf.train.piecewise_constant(global_step, cfg.train.lr_steps, cfg.train.learning_rate) optimizer = tf.train.AdamOptimizer(learning_rate=lr) # Calculate the gradients for each model tower. tower_grads = [] with tf.variable_scope(tf.get_variable_scope()): for i in range(cfg.train.num_gpus): with tf.device('/gpu:%d' % i): with tf.name_scope('%s_%d' % (cfg.train.tower, i)) as scope: model = BarDetNet(imgs_split[i], true_boxes_split[i], is_training) loss = model.compute_loss() tf.get_variable_scope().reuse_variables() grads = optimizer.compute_gradients(loss) tower_grads.append(grads) if i == 0: current_loss = loss update_op = tf.get_collection(tf.GraphKeys.UPDATE_OPS) vars_det = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope="BarDetNet") grads = average_gradients(tower_grads) with tf.control_dependencies(update_op): train_op = optimizer.minimize(loss, global_step=global_step, var_list=vars_det) # GPU config config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False) config.gpu_options.allow_growth = True sess = tf.Session(config=config) # Create a saver saver = tf.train.Saver() ckpt_dir = re.sub(r'examples/', '', cfg.ckpt_path_608) # init sess.run(tf.global_variables_initializer()) #saver.restore(sess, ckpt_dir+str(50000)+'_plate.ckpt-'+str(50000+1)) #gs = 50000 #sess.run(tf.assign(global_step, gs)) # running for i in range(0, cfg.train.max_batches): _, loss_ = sess.run([train_op, current_loss]) if(i % 1 == 0): print(i,': ', loss_) if i % 1000 == 0 and i < 10000: saver.save(sess, ckpt_dir+str(i)+'_charbar.ckpt', global_step=global_step, write_meta_graph=False) if i % 10000 == 0: saver.save(sess, ckpt_dir+str(i)+'_charbar.ckpt', global_step=global_step, write_meta_graph=False)
def train(): is_training = True # data pipeline imgs, true_boxes = gen_data_batch(re.sub(r'examples/', '', cfg.data_path), cfg.batch_size * cfg.train.num_gpus) imgs_split = tf.split(imgs, cfg.train.num_gpus) true_boxes_split = tf.split(true_boxes, cfg.train.num_gpus) global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0.), trainable=False) lr = tf.train.piecewise_constant(global_step, cfg.train.lr_steps, cfg.train.learning_rate) optimizer = tf.train.AdamOptimizer(learning_rate=lr) # Calculate the gradients for each model tower. tower_grads = [] summaries_buf = [] summaries = set() with tf.variable_scope(tf.get_variable_scope()): for i in range(cfg.train.num_gpus): with tf.device('/gpu:%d' % i): with tf.name_scope('%s_%d' % (cfg.train.tower, i)) as scope: model = PDetNet(imgs_split[i], true_boxes_split[i], is_training) loss = model.compute_loss() tf.get_variable_scope().reuse_variables() grads_and_vars = optimizer.compute_gradients(loss) # gradients_norm = summaries_gradients_norm(grads_and_vars) gradients_hist = summaries_gradients_hist(grads_and_vars) #summaries_buf.append(gradients_norm) summaries_buf.append(gradients_hist) ##sum_set = set() ##sum_set.add(tf.summary.scalar("loss", loss)) ##summaries_buf.append(sum_set) summaries_buf.append({tf.summary.scalar("loss", loss)}) # tower_grads.append(grads_and_vars) if i == 0: current_loss = loss update_op = tf.get_collection(tf.GraphKeys.UPDATE_OPS) vars_det = tf.get_collection( tf.GraphKeys.TRAINABLE_VARIABLES, scope="PDetNet") grads = average_gradients(tower_grads) with tf.control_dependencies(update_op): #train_op = optimizer.minimize(loss, global_step=global_step, var_list=vars_det) apply_gradient_op = optimizer.apply_gradients(grads, global_step=global_step) train_op = tf.group(apply_gradient_op, *update_op) # GPU config config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False) config.gpu_options.allow_growth = True sess = tf.Session(config=config) ##pruning add by lzlu # Parse pruning hyperparameters pruning_hparams = pruning.get_pruning_hparams().parse( cfg.prune.pruning_hparams) # Create a pruning object using the pruning hyperparameters pruning_obj = pruning.Pruning(pruning_hparams, global_step=global_step) # Use the pruning_obj to add ops to the training graph to update the masks # The conditional_mask_update_op will update the masks only when the # training step is in [begin_pruning_step, end_pruning_step] specified in # the pruning spec proto mask_update_op = pruning_obj.conditional_mask_update_op() # Use the pruning_obj to add summaries to the graph to track the sparsity # of each of the layers pruning_summaries = pruning_obj.add_pruning_summaries() summaries |= pruning_summaries for summ in summaries_buf: summaries |= summ summaries.add(tf.summary.scalar('lr', lr)) summary_op = tf.summary.merge(list(summaries), name='summary_op') if cfg.summary.summary_allowed: summary_writer = tf.summary.FileWriter( logdir=cfg.summary.logs_path, graph=sess.graph, flush_secs=cfg.summary.summary_secs) # Create a saver saver = tf.train.Saver() ckpt_dir = re.sub(r'examples/', '', cfg.ckpt_path_608) if cfg.train.fine_tune == 0: # init sess.run(tf.global_variables_initializer()) else: saver.restore(sess, cfg.train.rstd_path) # running for i in range(0, cfg.train.max_batches): _, loss_, gstep, sval, _ = sess.run( [train_op, current_loss, global_step, summary_op, mask_update_op]) if (i % 100 == 0): print(i, ': ', loss_) if i % 1000 == 0 and i < 10000: saver.save(sess, ckpt_dir + str(i) + '_plate.ckpt', global_step=global_step, write_meta_graph=False) if i % 10000 == 0: saver.save(sess, ckpt_dir + str(i) + '_plate.ckpt', global_step=global_step, write_meta_graph=False) if cfg.summary.summary_allowed and gstep % cfg.summary.summ_steps == 0: summary_writer.add_summary(sval, global_step=gstep)
#!/usr/bin/env python # encoding: utf-8 import numpy as np import tensorflow as tf import sys sys.path.append('..') from models.run_net import PDetNet from prepare_data.gen_data_batch import gen_data_batch from config import cfg import os os.environ["CUDA_VISIBLE_DEVICES"] = "0" cfg.train.image_resized = 608 imgs, true_boxes = gen_data_batch(cfg.data_path, cfg.batch_size) is_training = True model = PDetNet(imgs, true_boxes, is_training) loss = model.compute_loss() global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0.), trainable=False) lr = tf.train.piecewise_constant(global_step, cfg.train.lr_steps, cfg.train.learning_rate) optimizer = tf.train.AdamOptimizer(learning_rate=lr) update_op = tf.get_collection(tf.GraphKeys.UPDATE_OPS) vars_det = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope="PDetNet") with tf.control_dependencies(update_op): train_op = optimizer.minimize(loss,
def train(finetune): is_training = True # data pipeline imgs, true_boxes = gen_data_batch(cfg.data_path, cfg.batch_size*cfg.train.num_gpus) imgs_split = tf.split(imgs, cfg.train.num_gpus) true_boxes_split = tf.split(true_boxes, cfg.train.num_gpus) global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0.), trainable=False) #lr = tf.train.piecewise_constant(global_step, cfg.train.lr_steps, cfg.train.learning_rate) #optimizer = tf.train.AdamOptimizer(learning_rate=lr) learn_rate_decay_step = int(cfg.train.num_samples / cfg.batch_size / cfg.train.num_gpus * cfg.train.learn_rate_decay_epoch) learning_rate = tf.train.exponential_decay(cfg.train.learn_rate, global_step, learn_rate_decay_step, cfg.train.learn_rate_decay, staircase=True) #optimizer = tf.train.GradientDescentOptimizer(learning_rate) optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate) # Calculate the gradients for each model tower. tower_grads = [] with tf.variable_scope(tf.get_variable_scope()): for i in range(cfg.train.num_gpus): with tf.device('/gpu:%d' % i): with tf.name_scope('%s_%d' % (cfg.train.tower, i)) as scope: model = SenseClsNet(imgs_split[i], true_boxes_split[i], is_training) loss = model.compute_loss() tf.get_variable_scope().reuse_variables() grads = optimizer.compute_gradients(loss) tower_grads.append(grads) if i == 0: current_loss = loss update_op = tf.get_collection(tf.GraphKeys.UPDATE_OPS) # print(tf.GraphKeys.UPDATE_OPS) # print(update_op) # print(grads) # vars_det = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope="BioRecNet") grads = average_gradients(tower_grads) with tf.control_dependencies(update_op): apply_gradient_op = optimizer.apply_gradients(grads, global_step=global_step) train_op = tf.group(apply_gradient_op,*update_op) # GPU config config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False) config.gpu_options.allow_growth = True sess = tf.Session(config=config) # Create a saver saver = tf.train.Saver(max_to_keep=1000) ckpt_dir = cfg.ckpt_path if not os.path.exists(ckpt_dir): os.makedirs(ckpt_dir) # init sess.run(tf.global_variables_initializer()) if finetune: checkpoint = './pre_train.ckpt' # variables_to_restore = slim.get_variables_to_restore() # init_assign_op, init_feed_dict = slim.assign_from_checkpoint(checkpoint, variables_to_restore, ignore_missing_vars=True) # sess.run(init_assign_op, init_feed_dict) variables_to_restore = get_variables_to_restore(exclude_global_pool=True) init_assign_op, init_feed_dict = slim.assign_from_checkpoint(checkpoint, variables_to_restore, ignore_missing_vars=True) sess.run(init_assign_op, init_feed_dict) # running cnt_epoch = 0 for i in range(1, cfg.train.max_batches): _, loss_, lr_ = sess.run([train_op, current_loss, learning_rate]) if(i % 5 == 0): print(i,': ', loss_, ' lr: ', lr_) if int(i) % int(cfg.train.num_samples / cfg.train.num_gpus / cfg.batch_size) == 0: cnt_epoch += 1 saver.save(sess, ckpt_dir+'senceCls', global_step=cnt_epoch, write_meta_graph=True)
def train(finetune): is_training = True # data pipeline 批量引入数据 imgs, age_labels, age_vectors = gen_data_batch( cfg.data_path, cfg.batch_size * cfg.train.num_gpus) #配置函数 imgs = tf.reshape(imgs, (-1, imgs.get_shape()[2], imgs.get_shape()[3], imgs.get_shape()[4])) #将imgs处理成参数shape的列表形式 imgs_split = tf.split(imgs, cfg.train.num_gpus) #图片分割 age_labels_split = tf.split(age_labels, cfg.train.num_gpus) age_vectors_split = tf.split(age_vectors, cfg.train.num_gpus) global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0.), trainable=False) learn_rate_decay_step = int(cfg.train.num_samples / cfg.batch_size / cfg.train.num_gpus * cfg.train.learn_rate_decay_epoch) learning_rate = tf.train.exponential_decay(cfg.train.learn_rate, global_step, learn_rate_decay_step, cfg.train.learn_rate_decay, staircase=True) optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate) # Calculate the gradients for each model tower. 计算每个模型tower的梯度。 tower_grads = [] with tf.variable_scope(tf.get_variable_scope()): for i in range(cfg.train.num_gpus): with tf.device('/gpu:%d' % i): with tf.name_scope('%s_%d' % (cfg.train.tower, i)) as scope: model = C3AENet(imgs_split[i], age_labels_split[i], age_vectors_split[i], is_training) loss = model.compute_loss() tf.get_variable_scope().reuse_variables() grads = optimizer.compute_gradients(loss) tower_grads.append(grads) if i == 0: current_loss = loss update_op = tf.get_collection(tf.GraphKeys.UPDATE_OPS) grads = average_gradients(tower_grads) with tf.control_dependencies( update_op): #tf.control_dependencies是控制先update_op操作再计算的函数 apply_gradient_op = optimizer.apply_gradients(grads, global_step=global_step) train_op = tf.group(apply_gradient_op, *update_op) # GPU config config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False) config.gpu_options.allow_growth = True sess = tf.Session(config=config) # Create a saver 使用Saver类保存模型的参数 saver = tf.train.Saver(max_to_keep=1000) ckpt_dir = cfg.ckpt_path # init 初始化模型的参数 sess.run(tf.global_variables_initializer()) if finetune: checkpoint = './ckpt/pre_train.ckpt' variables_to_restore = get_variables_to_restore( exclude_global_pool=True) #variables_to_restore是TensorFlow为滑动平均值提供的一个函数,通过使用滑动平均值可以让神经网络模型更加的健壮 init_assign_op, init_feed_dict = slim.assign_from_checkpoint( checkpoint, variables_to_restore, ignore_missing_vars=True) sess.run(init_assign_op, init_feed_dict) # running cnt_epoch = 0 for i in range(1, cfg.train.max_batches): _, loss_, lr_ = sess.run([train_op, current_loss, learning_rate]) #训练节点、损失、学习率 进行训练 if (i % 10 == 0): print(i, ': ', loss_, ' lr: ', lr_) if int(i) % int(cfg.train.num_samples / cfg.train.num_gpus / cfg.batch_size) == 0: cnt_epoch += 1 saver.save(sess, ckpt_dir + 'C3AEDet', global_step=cnt_epoch, write_meta_graph=True)