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)