vectmap_node = tf.placeholder(tf.float32,
shape=(args.batchsize, output_h,
output_w, 38),
name='vectmap')
heatmap_node = tf.placeholder(tf.float32,
shape=(args.batchsize, output_h,
output_w, 19),
name='heatmap')
object_node = tf.placeholder(tf.float32,
shape=(args.batchsize, 2535, 85),
name='yolo_out')
# prepare data
if not args.remote_data:
df = get_dataflow_batch(args.datapath,
True,
args.batchsize,
img_path=args.imgpath)
print('train_df', df)
else:
# transfer inputs from ZMQ
df = RemoteDataZMQ(args.remote_data, hwm=3)
enqueuer = DataFlowToQueue(
df, [input_node, heatmap_node, vectmap_node, object_node],
queue_size=100)
q_inp, q_heat, q_vect, q_obj = enqueuer.dequeue()
print('inp/out', q_inp, q_heat, q_vect, q_obj)
df_valid = get_dataflow_batch(args.datapath,
False,
args.batchsize,
img_path=args.imgpath)
from tf_pose.pose_dataset import get_dataflow_batch
from tf_pose.pose_augment import set_network_input_wh, set_network_scale
if __name__ == '__main__':
"""
OpenPose Data Preparation might be a bottleneck for training.
You can run multiple workers to generate input batches in multi-nodes to make training process faster.
"""
parser = argparse.ArgumentParser(
description='Worker for preparing input batches.')
parser.add_argument('--datapath', type=str, default='/coco/annotations/')
parser.add_argument('--imgpath', type=str, default='/coco/')
parser.add_argument('--batchsize', type=int, default=64)
parser.add_argument('--train', type=bool, default=True)
parser.add_argument('--master',
type=str,
default='tcp://csi-cluster-gpu20.dakao.io:1027')
parser.add_argument('--input-width', type=int, default=368)
parser.add_argument('--input-height', type=int, default=368)
parser.add_argument('--scale-factor', type=int, default=2)
args = parser.parse_args()
set_network_input_wh(args.input_width, args.input_height)
set_network_scale(args.scale_factor)
df = get_dataflow_batch(args.datapath, args.train, args.batchsize,
args.imgpath)
send_dataflow_zmq(df, args.master, hwm=10)