def train(config):
if config['dataset'] == 'nyu':
dataset = NYUDataset(subset='testing', root_dir='/home/data/nyu/')
# (240, 180, 70), 6 * 14 = 84
actor_cnn_layer = (4, 8, 16, 32, 64, 128)
actor_fc_layer = (512, 512, 256)
critic_cnn_layer = (4, 8, 16, 32, 64, 128) # 768
critic_fc_layer = (512, 84, 512, 128)
elif config['dataset'] == 'icvl':
dataset = ICVLDataset(subset='testing',
root_dir='/hand_pose_data/icvl/')
# (140, 120, 60), 6 * 16 = 96
actor_cnn_layer = (4, 8, 16, 32, 64)
actor_fc_layer = (512, 512, 256)
critic_cnn_layer = (4, 8, 16, 32, 64) # 768
critic_fc_layer = (512, 96, 512, 128)
elif config['dataset'] == 'mrsa15':
# (180, 120, 70), 6 * 21 = 126
dataset = MRSADataset(subset='testing',
test_fold=config['mrsa_test_fold'],
root_dir='/hand_pose_data/mrsa15/')
actor_cnn_layer = (4, 8, 16, 64, 128, 256)
actor_fc_layer = (512, 512, 256)
critic_cnn_layer = (4, 8, 16, 64, 128, 256) # 512
critic_fc_layer = (512, 126, 512, 128)
else:
raise ValueError('Dataset name %s error...' % config['dataset'])
actor = Actor(scope='actor',
obs_dims=dataset.predefined_bbx,
ac_dim=6 * dataset.jnt_num,
cnn_layer=actor_cnn_layer,
fc_layer=actor_fc_layer,
tau=config['tau'],
lr=config['actor_lr'])
critic = Critic(scope='critic',
obs_dims=dataset.predefined_bbx,
ac_dim=6 * dataset.jnt_num,
cnn_layer=critic_cnn_layer,
fc_layer=critic_fc_layer,
tau=config['tau'],
lr=config['critic_lr'])
env = HandEnv(dataset_name=config['dataset'],
subset='training',
max_iters=config['max_iters'],
predefined_bbx=dataset.predefined_bbx)
sampler = Sampler(actor, critic, env, dataset, config['gamma'])
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
with tf.Session(config=tf_config) as sess:
sess.run(tf.global_variables_initializer())
root_dir = config['saved_model_path'] + '/' + config['dataset'] + '/'
# actor model
save_actor_dir = root_dir + 'actor.pkl'
if os.path.exists(save_actor_dir):
utils.loadFromFlat(actor.get_trainable_variables(), save_actor_dir)
print("Actor parameter loaded from %s" % save_actor_dir)
else:
ValueError('Actor parameter cannot be found at %s' %
save_actor_dir)
sess.run(actor.update_target_ops())
actor.load_sess(sess)
# critic model
save_critic_dir = root_dir + 'critic.pkl'
if os.path.exists(save_critic_dir):
utils.loadFromFlat(critic.get_trainable_variables(),
save_critic_dir)
print("Critic parameter loaded from %s" % save_critic_dir)
else:
ValueError('Critic parameter cannot be found at %s' %
save_critic_dir)
sess.run(critic.update_target_ops)
critic.load_sess(sess)
# testing
# list of tuples: (final_avg_distance, final_pose, filename, xyz_pose, depth_img, pose_bbx,
# coeff, normalized_rotate_pose, rotated_bbx)
results = sampler.test_multiple_samples()
predictions = []
for result in results:
pred_pose = transfer_pose(result, dataset.predefined_bbx)
# tuple: (filename, xyz_pose, pred_pose, depth_img)
predictions.append((result[2], result[3], pred_pose, result[4]))
with open(root_dir + 'predictions.pkl', 'wb') as f:
pickle.dump(predictions, f)
def test(config):
if config['dataset'] == 'nyu':
dataset = NYUDataset(subset='testing', root_dir='/hand_pose_data/nyu/')
elif config['dataset'] == 'icvl':
dataset = ICVLDataset(subset='testing', root_dir='/hand_pose_data/icvl/')
elif config['dataset'] == 'mrsa15':
dataset = MRSADataset(subset='testing', test_fold=config['mrsa_test_fold'],
root_dir='../../../hand_pose_data/mrsa15/')
else:
raise ValueError('Dataset name %s error...' % config['dataset'])
actor_root = Actor(scope='actor_root',
tau=config['tau'],
lr=config['learning_rate'],
obs_dims=config['root_obs_dims'],
cnn_layer=config['root_actor_cnn_layers'],
fc_layer=config['root_actor_fc_layers'])
critic_root = Critic(scope='critic_root',
tau=config['tau'],
lr=config['learning_rate'],
obs_dims=config['root_obs_dims'],
cnn_layer=config['root_critic_cnn_layers'],
fc_layer=config['root_critic_fc_layers'])
actor_chain = Actor(scope='actor_chain',
tau=config['tau'],
lr=config['learning_rate'],
obs_dims=config['chain_obs_dims'],
cnn_layer=config['chain_actor_cnn_layers'],
fc_layer=config['chain_actor_fc_layers'])
critic_chain = Critic(scope='critic_chain',
tau=config['tau'],
lr=config['learning_rate'],
obs_dims=config['chain_obs_dims'],
cnn_layer=config['chain_critic_cnn_layers'],
fc_layer=config['chain_critic_fc_layers'])
env = HandEnv(dataset=config['dataset'],
subset='training',
iter_per_joint=config['iter_per_joint'],
reward_beta=config['beta'])
sampler = Sampler(actor_root, critic_root, actor_chain, critic_chain, env, dataset)
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
with tf.Session(config=tf_config) as sess:
sess.run(tf.global_variables_initializer())
# load actor root model
root_dir = config['saved_model_path'] + '/' + config['dataset'] + '/'
saved_actor_dir = root_dir + config['actor_model_name'] + '_root.pkl'
if os.path.exists(saved_actor_dir):
utils.loadFromFlat(actor_root.get_trainable_variables(), saved_actor_dir)
print("Actor parameter loaded from %s" % saved_actor_dir)
else:
ValueError('Dir %s errors...' % saved_actor_dir)
actor_root.load_sess(sess)
# load critic root model
saved_critic_dir = root_dir + config['critic_model_name'] + '_root.pkl'
if os.path.exists(saved_critic_dir):
utils.loadFromFlat(critic_root.get_trainable_variables(), saved_critic_dir)
print("Critic parameter loaded from %s" % saved_critic_dir)
else:
ValueError('Dir %s errors...' % saved_critic_dir)
critic_root.load_sess(sess)
# load actor chain model
saved_actor_dir = root_dir + config['actor_model_name'] + '_chain.pkl'
if os.path.exists(saved_actor_dir):
utils.loadFromFlat(actor_chain.get_trainable_variables(), saved_actor_dir)
print("Actor_chain parameter loaded from %s" % saved_actor_dir)
else:
ValueError('Dir %s errors...' % saved_actor_dir)
actor_chain.load_sess(sess)
# load critic chain model
saved_critic_dir = root_dir + config['critic_model_name'] + '_chain.pkl'
if os.path.exists(saved_critic_dir):
utils.loadFromFlat(critic_chain.get_trainable_variables(), saved_critic_dir)
print("Critic_chain parameter loaded from %s" % saved_critic_dir)
else:
ValueError('Dir %s errors...' % saved_critic_dir)
critic_chain.load_sess(sess)
# test
results = sampler.test_multiple_samples()
def train_root(config):
if config['dataset'] == 'nyu':
dataset = NYUDataset(subset='training', root_dir='/home/data/nyu/')
elif config['dataset'] == 'icvl':
dataset = ICVLDataset(subset='training', root_dir='/hand_pose_data/icvl/')
elif config['dataset'] == 'mrsa15':
dataset = MRSADataset(subset='training', test_fold=config['mrsa_test_fold'],
root_dir='/hand_pose_data/mrsa15/')
else:
raise ValueError('Dataset name %s error...' % config['dataset'])
actor_root = Actor(scope='actor_root',
tau=config['tau'],
lr=config['learning_rate'],
obs_dims=config['root_obs_dims'],
cnn_layer=config['root_actor_cnn_layers'],
fc_layer=config['root_actor_fc_layers'])
critic_root = Critic(scope='critic_root',
tau=config['tau'],
lr=config['learning_rate'],
obs_dims=config['root_obs_dims'],
cnn_layer=config['root_critic_cnn_layers'],
fc_layer=config['root_critic_fc_layers'])
env = HandEnv(dataset=config['dataset'],
subset='training',
iter_per_joint=config['iter_per_joint'],
reward_beta=config['beta'])
root_buffer = ReplayBuffer(buffer_size=config['buffer_size'])
sampler = Sampler(actor_root, critic_root, None, None, env, dataset)
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
with tf.Session(config=tf_config) as sess:
sess.run(tf.global_variables_initializer())
# actor model
root_dir = config['saved_model_path'] + '/' + config['dataset'] + '/'
saved_actor_dir = root_dir + config['actor_model_name'] + '_root.pkl'
if os.path.exists(saved_actor_dir):
utils.loadFromFlat(actor_root.get_trainable_variables(), saved_actor_dir)
print("Actor parameter loaded from %s" % saved_actor_dir)
else:
print("[Warning]: initialize actor root model")
actor_root.load_sess(sess)
sess.run(actor_root.update_target_ops)
# critic model
saved_critic_dir = root_dir + config['critic_model_name'] + '_root.pkl'
if os.path.exists(saved_critic_dir):
utils.loadFromFlat(critic_root.get_trainable_variables(), saved_critic_dir)
print("Critic parameter loaded from %s" % saved_critic_dir)
else:
print("[Warning]: initialize critic root model")
critic_root.load_sess(sess)
sess.run(critic_root.update_target_ops)
i = 0
while i < config['n_rounds']:
i += 1
print('--------------------------------Round % i---------------------------------' % i)
# sampling
samples = sampler.collect_multiple_samples_root(config['files_per_time'])
root_buffer.add(samples)
for _ in range(config['n_iters']):
actor_loss_list, q_loss_list = [], []
for _ in range(config['update_iters']):
# get a mini-batch of data
state, action, reward, new_state, gamma = root_buffer.get_batch(config['batch_size'])
# update actor
q_gradient = critic_root.get_q_gradient(obs=state, ac=action)
_, actor_loss = actor_root.train(q_gradient=q_gradient[0], obs=state)
# update critic
next_ac = actor_root.get_target_action(obs=new_state)
_, q_loss = critic_root.train(obs=state, ac=action, next_obs=new_state,
next_ac=next_ac, r=reward, gamma=gamma)
actor_loss_list.append(actor_loss)
q_loss_list.append(q_loss)
# update target network
sess.run(actor_root.update_target_ops)
sess.run(critic_root.update_target_ops)
print('Actor average loss: {:.4f}, Critic: {:.4f}'
.format(np.mean(actor_loss_list), np.mean(q_loss_list)))
utils.saveToFlat(actor_root.get_trainable_variables(), saved_actor_dir)
utils.saveToFlat(critic_root.get_trainable_variables(), saved_critic_dir)
def test(config):
if config['dataset'] == 'nyu':
dataset = NYUDataset(subset='training', root_dir='/hand_pose_data/nyu/', predefined_bbx=(63, 63, 31))
# pre-trained model
pre_ac_dim = 3 * dataset.jnt_num
pre_cnn_layer = (8, 16, 32, 64, 128) # 512
pre_fc_layer = (512, 512, 256)
# actor-critic
ac_dim = 4 * (dataset.jnt_num - 1)
actor_cnn_layer = (8, 16, 32, 64, 128) # 512
actor_fc_layer = (512, 512, 256)
elif config['dataset'] == 'icvl':
dataset = ICVLDataset(subset='training', root_dir='/hand_pose_data/icvl/', predefined_bbx=(63, 63, 31))
# pre-trained model
pre_ac_dim = 3 * dataset.jnt_num
pre_cnn_layer = (8, 16, 32, 64, 128) # 512
pre_fc_layer = (512, 512, 256)
# actor-critic
ac_dim = 4 * (dataset.jnt_num - 1)
actor_cnn_layer = (8, 16, 32, 64, 128) # 512
actor_fc_layer = (512, 512, 256)
elif config['dataset'] == 'mrsa15':
# (180, 120, 70), 6 * 21 = 126
dataset = MRSADataset(subset='training', test_fold=config['mrsa_test_fold'],
root_dir='/hand_pose_data/mrsa15/', predefined_bbx=(63, 63, 31))
# pre-trained model
pre_ac_dim = 3 * dataset.jnt_num
pre_cnn_layer = (8, 16, 32, 64, 128) # 512
pre_fc_layer = (512, 512, 256)
# actor-critic
ac_dim = 4 * (dataset.jnt_num - 1)
actor_cnn_layer = (8, 16, 32, 64, 128) # 512
actor_fc_layer = (512, 512, 256)
else:
raise ValueError('Dataset name %s error...' % config['dataset'])
obs_dims = (dataset.predefined_bbx[2] + 1, dataset.predefined_bbx[1] + 1, dataset.predefined_bbx[0] + 1)
# build pretrain model
pretrain_model = Pretrain(scope='pretrain',
obs_dims=obs_dims + (1,),
cnn_layer=pre_cnn_layer,
fc_layer=pre_fc_layer,
ac_dim=pre_ac_dim)
# build actor
actor = Actor(scope='actor',
obs_dims=obs_dims + (2,),
ac_dim=ac_dim,
cnn_layer=actor_cnn_layer,
fc_layer=actor_fc_layer,
tau=config['tau'],
beta=config['beta'],
lr=config['actor_lr'])
# initialize environment
env = HandEnv(dataset_name=config['dataset'], subset='training', max_iters=config['max_iters'],
predefined_bbx=dataset.predefined_bbx, pretrained_model=pretrain_model,
reward_range=config['reward_range'], num_cpus=config['num_cpus'])
# initialize sampler
sampler = Sampler(actor, env, dataset, step_size=config['step_size'], gamma=config['gamma'])
buffer = ReplayBuffer(buffer_size=config['buffer_size'])
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
with tf.Session(config=tf_config) as sess:
sess.run(tf.global_variables_initializer())
root_dir = config['saved_model_path'] + config['dataset'] + '/'
# load pretrained model
if config['dataset'] == 'mrsa15':
model_save_dir = root_dir + config['dataset'] + '_' + config['mrsa_test_fold'] + '_pretrain.pkl'
else:
model_save_dir = root_dir + config['dataset'] + '_pretrain.pkl'
if os.path.exists(model_save_dir):
utils.loadFromFlat(pretrain_model.get_trainable_variables(), model_save_dir)
print("Pre-train parameter loaded from %s" % model_save_dir)
else:
raise ValueError('Model not found from %s' % model_save_dir)
# load actor model
save_actor_dir = root_dir + config['dataset'] + '_actor.pkl'
if os.path.exists(save_actor_dir):
utils.loadFromFlat(actor.get_trainable_variables(), save_actor_dir)
print("Actor parameter loaded from %s" % save_actor_dir)
else:
print("[Warning]: initialize the actor model")
actor.load_sess(sess)
def train(config):
if config['dataset'] == 'nyu':
dataset = NYUDataset(subset='training', root_dir='/hand_pose_data/nyu/', predefined_bbx=(63, 63, 31))
# pre-trained model
pre_ac_dim = 3 * dataset.jnt_num
pre_cnn_layer = (8, 16, 32, 64, 128) # 512
pre_fc_layer = (512, 512, 256)
# actor-critic
ac_dim = 4 * (dataset.jnt_num - 1)
actor_cnn_layer = (8, 16, 32, 64, 128) # 512
actor_fc_layer = (512, 512, 256)
critic_cnn_layer = (8, 16, 32, 64, 128) # 512
critic_fc_layer = (ac_dim, 512, 512, 128)
elif config['dataset'] == 'icvl':
dataset = ICVLDataset(subset='training', root_dir='/hand_pose_data/icvl/', predefined_bbx=(63, 63, 31))
# pre-trained model
pre_ac_dim = 3 * dataset.jnt_num
pre_cnn_layer = (8, 16, 32, 64, 128) # 512
pre_fc_layer = (512, 512, 256)
# actor-critic
ac_dim = 4 * (dataset.jnt_num - 1)
actor_cnn_layer = (8, 16, 32, 64, 128) # 512
actor_fc_layer = (512, 512, 256)
critic_cnn_layer = (8, 16, 32, 64, 128) # 512
critic_fc_layer = (ac_dim, 512, 512, 128)
elif config['dataset'] == 'mrsa15':
# (180, 120, 70), 6 * 21 = 126
dataset = MRSADataset(subset='training', test_fold=config['mrsa_test_fold'],
root_dir='/hand_pose_data/mrsa15/', predefined_bbx=(63, 63, 31))
# pre-trained model
pre_ac_dim = 3 * dataset.jnt_num
pre_cnn_layer = (8, 16, 32, 64, 128) # 512
pre_fc_layer = (512, 512, 256)
# actor-critic
ac_dim = 4 * (dataset.jnt_num - 1)
actor_cnn_layer = (8, 16, 32, 64, 128) # 512
actor_fc_layer = (512, 512, 256)
critic_cnn_layer = (8, 16, 32, 64, 128) # 512
critic_fc_layer = (ac_dim, 512, 512, 128)
else:
raise ValueError('Dataset name %s error...' % config['dataset'])
obs_dims = (dataset.predefined_bbx[2] + 1, dataset.predefined_bbx[1] + 1, dataset.predefined_bbx[0] + 1)
# build pretrain model
pretrain_model = Pretrain(scope='pretrain',
obs_dims=obs_dims+(1,),
cnn_layer=pre_cnn_layer,
fc_layer=pre_fc_layer,
ac_dim=pre_ac_dim)
# build actor and critic model
actor = Actor(scope='actor',
obs_dims=obs_dims+(2,),
ac_dim=ac_dim,
cnn_layer=actor_cnn_layer,
fc_layer=actor_fc_layer,
tau=config['tau'],
beta=config['beta'],
lr=config['actor_lr'])
critic = Critic(scope='critic',
obs_dims=obs_dims+(2,),
ac_dim=ac_dim,
cnn_layer=critic_cnn_layer,
fc_layer=critic_fc_layer,
tau=config['tau'],
lr=config['critic_lr'])
# initialize environment
env = HandEnv(dataset_name=config['dataset'], subset='training', max_iters=config['max_iters'],
predefined_bbx=dataset.predefined_bbx, pretrained_model=pretrain_model,
reward_range=config['reward_range'], num_cpus=config['num_cpus'])
# initialize sampler
sampler = Sampler(actor, env, dataset, step_size=config['step_size'], gamma=config['gamma'])
buffer = ReplayBuffer(buffer_size=config['buffer_size'])
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
with tf.Session(config=tf_config) as sess:
sess.run(tf.global_variables_initializer())
root_dir = config['saved_model_path'] + config['dataset'] + '/'
writer = SummaryWriter(log_dir=root_dir)
# load pretrained model
if config['dataset'] == 'mrsa15':
model_save_dir = root_dir + config['dataset'] + '_' + config['mrsa_test_fold'] + '_pretrain.pkl'
else:
model_save_dir = root_dir + config['dataset'] + '_pretrain.pkl'
if os.path.exists(model_save_dir):
utils.loadFromFlat(pretrain_model.get_trainable_variables(), model_save_dir)
print("Pre-train parameter loaded from %s" % model_save_dir)
else:
raise ValueError('Model not found from %s' % model_save_dir)
# load actor model
save_actor_dir = root_dir + config['dataset'] + '_actor.pkl'
if os.path.exists(save_actor_dir):
# utils.loadFromFlat(actor.get_trainable_variables(), save_actor_dir)
print("Actor parameter loaded from %s" % save_actor_dir)
else:
print("[Warning]: initialize the actor model")
sess.run(actor.update_target_ops)
actor.load_sess(sess)
# critic model
save_critic_dir = root_dir + config['dataset'] + '_critic.pkl'
if os.path.exists(save_critic_dir):
# utils.loadFromFlat(critic.get_trainable_variables(), save_critic_dir)
print("Critic parameter loaded from %s" % save_critic_dir)
else:
print("[Warning]: initialize critic root model")
sess.run(critic.update_target_ops)
critic.load_sess(sess)
best_max_error = 20
test_examples = sampler.aggregate_test_samples()
for i in range(config['n_rounds']):
print('--------------------------------Round % i---------------------------------' % i)
if i % config['test_gap'] == 0:
# test
start_time = time.time()
print('>>>number of examples for testing: %i(%i)'
% (min(2*config['num_batch_samples'], len(test_examples)), len(test_examples)))
examples = random.sample(test_examples, min(2*config['num_batch_samples'], len(test_examples)))
max_error, rs = sampler.test_batch_samples(examples, 8*config['batch_size'], sess)
writer.add_histogram('RL_' + config['dataset'] + '_final_rewards', rs, i)
writer.add_histogram('RL_' + config['dataset'] + '_max_error', max_error, i)
writer.add_scalar('RL_' + config['dataset'] + '_mean_max_error', np.mean(max_error), i)
if best_max_error > np.mean(max_error):
# save model
utils.saveToFlat(actor.get_trainable_variables(), save_actor_dir)
utils.saveToFlat(critic.get_trainable_variables(), save_critic_dir)
best_max_error = np.mean(max_error)
print('>>>Model save as %s' % save_actor_dir)
end_time = time.time()
print('>>>Testing: Average max error {:.2f}, average reward {:.2f}, time used {:.2f}s'
.format(np.mean(max_error), np.mean(rs), end_time-start_time))
start_time = time.time()
# sampling
experiences, rs = sampler.collect_experiences(num_files=config['files_per_time'],
num_batch_samples=config['num_batch_samples'],
batch_size=8*config['batch_size'],
sess=sess,
num_cpus=config['num_cpus'])
buffer.add(experiences)
end_time = time.time()
print('Sampling: time used %.2fs, buffer size %i' % (end_time-start_time, buffer.count()))
# training
start_time = time.time()
actor_loss_list, q_loss_list = [], []
for _ in range(config['train_iters']):
# get a mini-batch of data
action, reward, gamma, state, new_state = buffer.get_batch(config['batch_size'])
# update actor
q_gradient = critic.get_q_gradient(obs=state, ac=action, dropout_prob=1.0)
_, actor_loss, global_step, actor_acs = \
actor.train(q_gradient=q_gradient[0], obs=state, dropout_prob=0.5, step_size=config['step_size'])
# update critic
next_ac = actor.get_target_action(obs=new_state, dropout_prob=1.0, step_size=config['step_size'])
_, critic_loss = critic.train(obs=state, ac=action, next_obs=new_state, next_ac=next_ac,
r=reward, gamma=gamma, dropout_prob=0.5)
# record result
actor_loss_list.append(np.mean(actor_loss))
q_loss_list.append(critic_loss)
writer.add_scalar('RL_' + config['dataset'] + '_actor_loss', np.mean(actor_loss), global_step)
writer.add_scalar('RL_' + config['dataset'] + '_critic_loss', critic_loss, global_step)
if global_step % config['update_iters'] == 0:
# update target network
sess.run([actor.update_target_ops, critic.update_target_ops])
print('Average loss: actor {:.4f}, critic: {:.4f}, training steps: {}, '
'average acs {:.4f}, average q-gradients {:.4f}'
.format(np.mean(actor_loss_list), np.mean(q_loss_list), global_step,
np.mean(actor_acs), np.mean(q_gradient)))
end_time = time.time()
print('Training time used: {:.2f}s, training steps: {}'.format(end_time - start_time, global_step))
writer.close()
def pre_test(config):
if config['dataset'] == 'nyu':
dataset = NYUDataset(subset='training',
root_dir='/hand_pose_data/nyu/',
predefined_bbx=(63, 63, 31))
ac_dim = 3 * dataset.jnt_num
# (160, 120, 70), 6 * 14 = 84
cnn_layer = (8, 16, 32, 64, 128) # 512
fc_layer = (512, 512, 256)
elif config['dataset'] == 'icvl':
dataset = ICVLDataset(subset='training',
root_dir='/hand_pose_data/icvl/',
predefined_bbx=(63, 63, 31))
ac_dim = 3 * dataset.jnt_num
# (140, 120, 60), 6 * 16 = 96
cnn_layer = (8, 16, 32, 64, 128) # 512
fc_layer = (512, 512, 256)
elif config['dataset'] == 'mrsa15':
# (180, 120, 70), 6 * 21 = 126
dataset = MRSADataset(subset='training',
test_fold=config['mrsa_test_fold'],
root_dir='/hand_pose_data/mrsa15/',
predefined_bbx=(63, 63, 31))
ac_dim = 3 * dataset.jnt_num
cnn_layer = (8, 16, 32, 64, 128) # 512
fc_layer = (512, 512, 256)
else:
raise ValueError('Dataset name %s error...' % config['dataset'])
obs_dims = (dataset.predefined_bbx[2] + 1, dataset.predefined_bbx[1] + 1,
dataset.predefined_bbx[0] + 1, 1)
env = HandEnv(dataset_name=config['dataset'],
subset='training',
max_iters=5,
predefined_bbx=dataset.predefined_bbx,
pretrained_model=None)
scope = 'pre_train'
batch_size = config['batch_size']
tf.reset_default_graph()
model = Pretrain(scope, obs_dims, cnn_layer, fc_layer,
ac_dim) # model.obs, model.ac, model.dropout_prob
tf_label = tf.placeholder(shape=(None, ac_dim),
dtype=tf.float32,
name='action')
tf_max_error = tf.sqrt(
tf.reduce_max(tf.reduce_sum(tf.reshape(tf.square(model.ac - tf_label),
[-1, int(ac_dim / 3), 3]),
axis=2),
axis=1))
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
with tf.Session(config=tf_config) as sess:
sess.run(tf.global_variables_initializer())
root_dir = config['saved_model_path'] + config['dataset'] + '/'
if config['dataset'] == 'mrsa15':
model_save_dir = root_dir + config['dataset'] + '_' + config[
'mrsa_test_fold'] + '_pretrain.pkl'
else:
model_save_dir = root_dir + config['dataset'] + '_pretrain.pkl'
if os.path.exists(model_save_dir):
utils.loadFromFlat(model.get_trainable_variables(), model_save_dir)
print("Pre-train parameter loaded from %s" % model_save_dir)
else:
raise ValueError('Model not found from %s' % model_save_dir)
x_test, y_test, other_data = collect_test_samples(
env, dataset, config['num_cpus'], ac_dim, False)
# # check function transfer_pretest_pose()
# i = 0
# filename, pred_pose = transfer_pretest_pose(y_test[i, :].reshape([int(ac_dim/3), 3]), other_data[i],
# dataset.jnt_num, env.chains_idx,
# dataset.predefined_bbx, dataset.camera_cfg, config['dataset'])
# print(filename, '\n', pred_pose)
n_test = x_test.shape[0]
start_time = time.time()
pred_list, max_error_list = [], []
for j in range(n_test // batch_size + 1):
idx1 = j * batch_size
idx2 = min((j + 1) * batch_size, n_test)
preds, max_error = sess.run(
[model.ac, tf_max_error],
feed_dict={
model.obs: x_test[idx1:idx2, ...],
model.dropout_prob: 1.0,
tf_label: y_test[idx1:idx2, ...]
})
pred_list.append(preds)
max_error_list.append(max_error)
pred_poses = np.concatenate(pred_list,
axis=0).reshape([-1,
int(ac_dim / 3), 3])
max_error_list = np.hstack(max_error_list)
print(
'Prediction of initial pose is done...({} examples, mean_max_error {:.3f})'
.format(pred_poses.shape[0], np.mean(max_error_list)))
results = []
pool = multiprocessing.Pool(config['num_cpus'])
for i in range(pred_poses.shape[0]):
results.append(
pool.apply_async(transfer_pretest_pose, (
pred_poses[i, :],
other_data[i],
dataset.jnt_num,
env.chains_idx,
dataset.predefined_bbx,
dataset.camera_cfg,
config['dataset'],
)))
pool.close()
pool.join()
preds = {}
for result in results:
_filename, pred_pose = result.get()
if config['dataset'] == 'nyu':
# depth_1_0000601.png, test/depth_1_0000601.png
filename = 'test/' + _filename
elif config['dataset'] == 'mrsa15':
# /hand_pose_data/mrsa15/dataset/P8/I/000307_depth.bin, P8/I/000307_depth.bin
filename = '/'.join(_filename.split('/')[-3:])
elif config['dataset'] == 'icvl':
# 'test_seq_1/image_0002.png'
filename = _filename
preds[filename] = pred_pose
end_time = time.time()
print('Pretest finished. ({:.2f}s)'.format(end_time - start_time))
if config['dataset'] == 'mrsa15':
# save predictions
with open(
root_dir +
'pre_test_results_%s.pkl' % config['mrsa_test_fold'],
'wb') as f:
pickle.dump(preds, f)
print('Predictions are saved as %s.' %
(root_dir +
'pre_test_results_%s.pkl' % config['mrsa_test_fold']))
else:
# save predictions
with open(root_dir + 'pre_test_results.pkl', 'wb') as f:
pickle.dump(preds, f)
print('Predictions are saved as %s.' % root_dir +
'pre_test_results.pkl')
# write text file
write_text_file(
preds, config['dataset'], root_dir + config['dataset'] +
'_pretrain_rl_pose_estimation.txt')
def pre_train(config):
if config['dataset'] == 'nyu':
dataset = NYUDataset(subset='training',
root_dir='/hand_pose_data/nyu/',
predefined_bbx=(63, 63, 31))
ac_dim = 3 * dataset.jnt_num
weights = np.ones([1, dataset.jnt_num])
weights[0, 13] = config['root_weight'] # weight root joint error
cnn_layer = (8, 16, 32, 64, 128) # 512
fc_layer = (512, 512, 256)
elif config['dataset'] == 'icvl':
dataset = ICVLDataset(subset='training',
root_dir='/hand_pose_data/icvl/',
predefined_bbx=(63, 63, 31))
ac_dim = 3 * dataset.jnt_num
weights = np.ones([1, dataset.jnt_num])
weights[0, 0] = config['root_weight'] # weight root joint error
cnn_layer = (8, 16, 32, 64, 128) # 512
fc_layer = (512, 512, 256)
elif config['dataset'] == 'mrsa15':
# (180, 120, 70), 6 * 21 = 126
dataset = MRSADataset(subset='training',
test_fold=config['mrsa_test_fold'],
root_dir='/hand_pose_data/mrsa15/',
predefined_bbx=(63, 63, 31))
ac_dim = 3 * dataset.jnt_num
weights = np.ones([1, dataset.jnt_num])
weights[0, 0] = config['root_weight'] # weight root joint error
cnn_layer = (8, 16, 32, 64, 128) # 512
fc_layer = (512, 512, 256)
else:
raise ValueError('Dataset name %s error...' % config['dataset'])
print('Loss Weights:', weights)
obs_dims = (dataset.predefined_bbx[2] + 1, dataset.predefined_bbx[1] + 1,
dataset.predefined_bbx[0] + 1, 1)
env = HandEnv(dataset_name=config['dataset'],
subset='training',
max_iters=5,
predefined_bbx=dataset.predefined_bbx,
pretrained_model=None)
scope = 'pre_train'
batch_size = config['batch_size']
# define model and loss
model = Pretrain(scope, obs_dims, cnn_layer, fc_layer,
ac_dim) # model.obs, model.ac, model.dropout_prob
tf_label = tf.placeholder(shape=(None, ac_dim),
dtype=tf.float32,
name='action')
tf_weights = tf.placeholder(shape=(1, dataset.jnt_num),
dtype=tf.float32,
name='action')
# average joint mse error
tf_mse = tf.reduce_mean(tf_weights * tf.reduce_sum(tf.reshape(
tf.square(model.ac - tf_label), [-1, int(ac_dim / 3), 3]),
axis=2),
axis=1)
tf_loss = tf.reduce_mean(tf_mse) # average over mini-batch
tf_max_error = tf.sqrt(
tf.reduce_max(tf.reduce_sum(tf.reshape(tf.square(model.ac - tf_label),
[-1, int(ac_dim / 3), 3]),
axis=2),
axis=1))
global_step = tf.Variable(0, trainable=False, name='step')
lr = tf.train.exponential_decay(config['lr_start'], global_step,
config['lr_decay_iters'],
config['lr_decay_rate'])
optimizer = tf.train.AdamOptimizer(learning_rate=lr).minimize(
loss=tf_loss, global_step=global_step)
# optimizer = tf.train.RMSPropOptimizer(learning_rate=lr).minimize(loss=tf_loss, global_step=global_step)
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
with tf.Session(config=tf_config) as sess:
sess.run(tf.global_variables_initializer())
root_dir = config['saved_model_path'] + config['dataset'] + '/'
writer = SummaryWriter(log_dir=root_dir)
if config['dataset'] == 'mrsa15':
model_save_dir = root_dir + config['dataset'] + '_' + config[
'mrsa_test_fold'] + '_pretrain.pkl'
else:
model_save_dir = root_dir + config['dataset'] + '_pretrain.pkl'
if os.path.exists(model_save_dir) and not config['new_training']:
utils.loadFromFlat(model.get_trainable_variables(), model_save_dir)
print("Pre-train parameter loaded from %s" % model_save_dir)
else:
print("[Warning]: initialize the pre-train model")
x_test, y_test, _ = collect_test_samples(env, dataset,
config['num_cpus'], ac_dim)
n_test = x_test.shape[0]
print('test samples %i' % n_test)
best_loss = 1000
i = 0
while i < config['n_rounds']:
i += 1
print(
'--------------------------------Round % i---------------------------------'
% i)
# test
loss_list, max_error_list = [], []
if i % config['test_gap'] == 1:
start_time = time.time()
for j in range(n_test // batch_size + 1):
idx1 = j * batch_size
idx2 = min((j + 1) * batch_size, n_test)
batch_loss, batch_max_error = sess.run(
[tf_mse, tf_max_error],
feed_dict={
model.obs: x_test[idx1:idx2, ...],
tf_label: y_test[idx1:idx2],
model.dropout_prob: 1.0,
tf_weights: weights
})
loss_list.append(batch_loss)
max_error_list.append(batch_max_error)
test_loss = np.mean(np.hstack(loss_list))
max_error = np.hstack(max_error_list)
writer.add_scalar(config['dataset'] + '_test_loss', test_loss,
i)
writer.add_histogram(config['dataset'] + '_max_error',
max_error, i)
end_time = time.time()
print(
'>>> Testing loss: {:.4f}, best loss {:.4f}, mean_max_error {:.4f}， time used: {:.2f}s'
.format(test_loss, best_loss, np.mean(max_error),
end_time - start_time))
if best_loss > test_loss:
utils.saveToFlat(model.get_trainable_variables(),
model_save_dir)
best_loss = test_loss.copy()
print('>>> Model saved... best loss {:.4f}'.format(
best_loss))
# train
start_time = time.time()
x_train, y_train = collect_train_samples(
env, dataset, config['files_per_time'],
config['samples_per_time'], config['num_cpus'], ac_dim)
print('Collected samples {}'.format(x_train.shape[0]))
loss_list = []
for _ in range(config['train_iters']):
batch_idx = np.random.randint(0, x_train.shape[0], batch_size)
_, batch_loss, step = sess.run(
[optimizer, tf_loss, global_step],
feed_dict={
model.obs: x_train[batch_idx, ...],
tf_label: y_train[batch_idx],
model.dropout_prob: 0.5,
tf_weights: weights
})
loss_list.append(batch_loss)
end_time = time.time()
writer.add_scalar(config['dataset'] + '_train_loss',
np.mean(loss_list), i)
print(
'Training loss: {:.4f}, time used: {:.2f}s, step: {:d}'.format(
np.mean(loss_list), end_time - start_time, step))
writer.close()