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() 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])