def setup_pose_prediction(cfg):
# 该函数返回session，输入算子，输出算子
# 输入的placeholder
inputs = tf.placeholder(tf.float32, shape=[cfg.batch_size , None, None, 3])
# 获得测试的网络
net_heads = pose_net(cfg).test(inputs)
# 获得网络输出的heatmap
outputs = [net_heads['part_prob']]
if cfg.location_refinement:
# 是否使用了关节位置精细化定位
outputs.append(net_heads['locref'])
# 定义一个恢复权重的算子
restorer = tf.train.Saver()
sess = tf.Session()
# 初始化全局核局部变量
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Restore variables from disk.
# 从cfg.init_weights位置恢复权重
restorer.restore(sess, cfg.init_weights)
# 返回session，输入算子，输出算子
return sess, inputs, outputs
def train():
setup_logging()
cfg = load_config()
dataset = create_dataset(cfg)
batch_spec = get_batch_spec(cfg)
batch, enqueue_op, placeholders = setup_preloading(batch_spec)
losses = pose_net(cfg).train(batch)
total_loss = losses['total_loss']
print("total_loss:", total_loss)
for k, t in losses.items():
tf.summary.scalar(k, t)
merged_summaries = tf.summary.merge_all()
print("merged_summaries:", merged_summaries)
variables_to_restore = slim.get_variables_to_restore(include=["resnet_v1"])
restorer = tf.train.Saver(variables_to_restore)
saver = tf.train.Saver(max_to_keep=5)
sess = tf.Session()
coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders)
train_writer = tf.summary.FileWriter(cfg.log_dir, sess.graph)
learning_rate, train_op = get_optimizer(total_loss, cfg)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Restore variables from disk.
restorer.restore(sess, cfg.init_weights)
max_iter = int(cfg.multi_step[-1][1])
display_iters = cfg.display_iters
cum_loss = 0.0
lr_gen = LearningRate(cfg)
for it in range(max_iter + 1):
current_lr = lr_gen.get_lr(it)
[_, loss_val,
summary] = sess.run([train_op, total_loss, merged_summaries],
feed_dict={learning_rate: current_lr})
cum_loss += loss_val
train_writer.add_summary(summary, it)
if it % display_iters == 0:
average_loss = cum_loss / display_iters
cum_loss = 0.0
logging.info("iteration: {} loss: {} lr: {}".format(
it, "{0:.4f}".format(average_loss), current_lr))
# Save snapshot
if (it % cfg.save_iters == 0 and it != 0) or it == max_iter:
model_name = cfg.snapshot_prefix
saver.save(sess, model_name, global_step=it)
sess.close()
coord.request_stop()
coord.join([thread])
def setup_pose_prediction(cfg):
inputs = tf.placeholder(tf.float32, shape=[cfg.batch_size, None, None, 3])
outputs = pose_net(cfg).test(inputs)
restorer = tf.train.Saver()
sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Restore variables from disk.
restorer.restore(sess, cfg.init_weights)
return sess, inputs, outputs
def setup_pose_prediction(cfg):
inputs = tf.placeholder(tf.float32, shape=[cfg.batch_size, None, None, 3])
outputs = pose_net(cfg).test(inputs)
restorer = tf.train.Saver()
sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Restore variables from disk.
restorer.restore(sess, cfg.init_weights)
return sess, inputs, outputs
def setup_pose_prediction(cfg):
inputs = tf.placeholder(tf.float32, shape=[cfg.batch_size, None, None, 3])
outputs = pose_net(cfg).test(inputs)
restorer = tf.train.Saver()
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=cfg.per_process_gpu_memory_fraction)
config = tf.ConfigProto(gpu_options=gpu_options)
sess = tf.Session(config=config)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Restore variables from disk.
restorer.restore(sess, cfg.init_weights)
return sess, inputs, outputs
def setup_pose_prediction(cfg):
inputs = tf.placeholder(tf.float32, shape=[cfg.batch_size, None, None, 3])
net_heads = pose_net(cfg).test(inputs)
outputs = [net_heads['part_prob']]
if cfg.location_refinement:
outputs.append(net_heads['locref'])
restorer = tf.train.Saver()
sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Restore variables from disk.
restorer.restore(sess, cfg.init_weights)
return sess, inputs, outputs
def setup_pose_prediction(cfg, init_weights):
inputs = tf.placeholder(tf.float32, shape=[cfg.batch_size , None, None, 3])
net_heads = pose_net(cfg).test(inputs)
outputs = [net_heads['part_prob']]
if cfg.location_refinement:
outputs.append(net_heads['locref'])
restorer = tf.train.Saver()
sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Restore variables from disk.
restorer.restore(sess, init_weights)
return sess, inputs, outputs
def setup_pose_prediction(cfg):
inputs = tf.placeholder(tf.float32, shape=[cfg.batch_size, None, None, 3])
outputs = pose_net(cfg).test(inputs)
restorer = tf.train.Saver()
sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
#create tensorboard graph file during prediction#
writer = tf.summary.FileWriter("testboard", sess.graph)
writer.close()
# Restore variables from disk.
restorer.restore(sess, cfg.init_weights)
return sess, inputs, outputs
def train():
# 设置日志
setup_logging()
# 载入训练配置文件pose_cfg.yaml
cfg = load_config()
# 创建数据集类的实例
dataset = create_dataset(cfg)
# 获取batch_spec
# 包含输入图片大小
# 关节heatmap的大小
# 关节weight的大小
# 精细化heatmap的大小
# 精细化mask的大小
batch_spec = get_batch_spec(cfg)
# 根据batch_spec产生入队操作、placeholder和batch数据
batch, enqueue_op, placeholders = setup_preloading(batch_spec)
# 生成网络结构并且产生losses op
# 其中losses包括很多类型的loss
losses = pose_net(cfg).train(batch)
total_loss = losses['total_loss']
# 把多个loss合并起来
for k, t in losses.items():
# return a scalar Tensor of type string which contains a Summary protobuf.
tf.summary.scalar(k, t)
# returns a scalar Tensor of type string containing the serialized
# Summary protocol buffer resulting from the merging
merged_summaries = tf.summary.merge_all()
# 获取/resnet_v1下面的所有的变量
variables_to_restore = slim.get_variables_to_restore(include=["resnet_v1"])
# Create the saver which will be used to restore the variables.
# 创建一个恢复resent_v1的权重的op
restorer = tf.train.Saver(variables_to_restore)
# 创建一个保存训练状态的op
saver = tf.train.Saver(max_to_keep=5)
sess = tf.Session()
# 开启一个线程去读取数据并且装入到队列
coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders)
# 打开一个训练的记录器
train_writer = tf.summary.FileWriter(cfg.log_dir, sess.graph)
# 获取train_op和学习率op
learning_rate, train_op = get_optimizer(total_loss, cfg)
# 初始化全局和局部变量
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Restore variables from disk.从文件中读取权重到内存
restorer.restore(sess, cfg.init_weights)
# 从配置文件获取最大迭代次数
max_iter = int(cfg.multi_step[-1][1])
display_iters = cfg.display_iters
cum_loss = 0.0
# 生成一个学习率产生器的实例
lr_gen = LearningRate(cfg)
for it in range(max_iter + 1):
# 根据当前迭代的次数产生一个学习率
current_lr = lr_gen.get_lr(it)
# 进行训练
[_, loss_val,
summary] = sess.run([train_op, total_loss, merged_summaries],
feed_dict={learning_rate: current_lr})
# 累加loss
cum_loss += loss_val
# 将迭代次数保存起来
train_writer.add_summary(summary, it)
if it % display_iters == 0: # 每隔display_iters就显示一次 loss
average_loss = cum_loss / display_iters # 平均loss
cum_loss = 0.0
logging.info("iteration: {} loss: {} lr: {}".format(
it, "{0:.4f}".format(average_loss), current_lr))
# Save snapshot
# 每隔cfg.save_iters次就会保存
if (it % cfg.save_iters == 0 and it != 0) or it == max_iter:
# 获得模型的名称
model_name = cfg.snapshot_prefix
# 保存模型
saver.save(sess, model_name, global_step=it)
sess.close()
# 请求数据读取线程停止
coord.request_stop()
# 等待数据读取线程结束
coord.join([thread])
def train():
setup_logging()
cfg = load_config()
dataset = create_dataset(cfg)
batch_spec = get_batch_spec(cfg)
batch, enqueue_op, placeholders = setup_preloading(batch_spec)
losses = pose_net(cfg).train(batch)
total_loss = losses['total_loss']
for k, t in losses.items():
tf.summary.scalar(k, t)
merged_summaries = tf.summary.merge_all()
variables_to_restore = slim.get_variables_to_restore(include=["resnet_v1"])
restorer = tf.train.Saver(variables_to_restore)
saver = tf.train.Saver(max_to_keep=5)
sess = tf.Session()
coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders)
train_writer = tf.summary.FileWriter(cfg.log_dir, sess.graph)
learning_rate, train_op = get_optimizer(total_loss, cfg)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Restore variables from disk.
restorer.restore(sess, cfg.init_weights)
max_iter = int(cfg.multi_step[-1][1])
display_iters = cfg.display_iters
cum_loss = 0.0
lr_gen = LearningRate(cfg)
for it in range(max_iter+1):
current_lr = lr_gen.get_lr(it)
[_, loss_val, summary] = sess.run([train_op, total_loss, merged_summaries],
feed_dict={learning_rate: current_lr})
cum_loss += loss_val
train_writer.add_summary(summary, it)
if it % display_iters == 0:
average_loss = cum_loss / display_iters
cum_loss = 0.0
logging.info("iteration: {} loss: {} lr: {}"
.format(it, "{0:.4f}".format(average_loss), current_lr))
# Save snapshot
if (it % cfg.save_iters == 0 and it != 0) or it == max_iter:
model_name = cfg.snapshot_prefix
saver.save(sess, model_name, global_step=it)
sess.close()
coord.request_stop()
coord.join([thread])
def train():
setup_logging()
cfg = load_config()
# load newest snapshot
snapshots = [fn.split('.')[0] for fn in os.listdir(os.getcwd()) if "index" in fn]
if len(snapshots) > 0:
iters = np.array([int(fn.split('-')[1]) for fn in snapshots])
cfg['init_weights'] = snapshots[iters.argmax()]
start = iters.max()
else:
start = 0
dataset = create_dataset(cfg)
batch_spec = get_batch_spec(cfg)
batch, enqueue_op, placeholders = setup_preloading(batch_spec)
losses = pose_net(cfg).train(batch)
total_loss = losses['total_loss']
for k, t in losses.items():
tf.summary.scalar(k, t)
merged_summaries = tf.summary.merge_all()
if start==0:
variables_to_restore = slim.get_variables_to_restore(include=["resnet_v1"])
restorer = tf.train.Saver(variables_to_restore)
else:
restorer = tf.train.Saver()
saver = tf.train.Saver(max_to_keep=5)
sess = tf.Session()
coord, thread = start_preloading(sess, enqueue_op, dataset, placeholders)
train_writer = tf.summary.FileWriter(cfg.log_dir, sess.graph)
learning_rate, train_op = get_optimizer(total_loss, cfg)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
# Restore variables from disk.
restorer.restore(sess, cfg.init_weights)
max_iter = int(cfg.multi_step[-1][1])
display_iters = cfg.display_iters
cum_loss = 0.0
lr_gen = LearningRate(cfg, start)
startTime = time.time()
for it in range(start, max_iter+1):
current_lr = lr_gen.get_lr(it)
[_, loss_val, summary] = sess.run([train_op, total_loss, merged_summaries],
feed_dict={learning_rate: current_lr})
cum_loss += loss_val
train_writer.add_summary(summary, it)
if it % display_iters == 0:
average_loss = cum_loss / display_iters
cum_loss = 0.0
elapsed = timedelta(seconds=(time.time()-startTime))
logging.info("iteration: {} loss: {} lr: {} time: {}"
.format(it, "{0:.4f}".format(average_loss), current_lr, elapsed))
# Save snapshot
if (it % cfg.save_iters == 0 and it != start) or it == max_iter:
model_name = cfg.snapshot_prefix
saver.save(sess, model_name, global_step=it)
sess.close()
coord.request_stop()
coord.join([thread])
