def write_tf():
data = ut.load('../results/grasp-dset/v5/train.pk')
tf_file = '../results/grasp-dset/v5/train.tf'
assert not os.path.exists(tf_file)
writer = tf.python_io.TFRecordWriter(tf_file)
for d in ut.time_est(data):
fbl = lambda x: tf.train.Feature(bytes_list=tf.train.BytesList(value=
[x]))
fl = lambda x: tf.train.Feature(float_list=tf.train.FloatList(
value=map(float, x.flatten())))
il = lambda x: tf.train.Feature(int64_list=tf.train.Int64List(value=x))
feat = {
'gel0_pre': fbl(d['gel0_pre']),
'gel1_pre': fbl(d['gel1_pre']),
'gel0_post': fbl(d['gel0_post']),
'gel1_post': fbl(d['gel1_post']),
'im0_pre': fbl(d['im0_pre']),
'im0_post': fbl(d['im0_post']),
'im1_pre': fbl(d['im1_pre']),
'im1_post': fbl(d['im1_post']),
'depth0_pre': fl(d['depth0_pre']),
'depth0_post': fl(d['depth0_post']),
'end_effector': fl(d['end_effector']),
'initial_press_prob': fl(d['initial_press_prob']),
'is_gripping': il([d['is_gripping']])
}
ex = tf.train.Example(features=tf.train.Features(feature=feat))
writer.write(ex.SerializeToString())
writer.close()
full_dur = arg.duration
#full_dur = min(arg.duration, ut.video_length(arg.vid_file))
#full_dur = arg.duration
step_dur = arg.clip_dur / 2.
filled = np.zeros(int(np.ceil(full_dur * pr.samp_sr)), 'bool')
full_samples_fg = np.zeros(filled.shape, 'float32')
full_samples_bg = np.zeros(filled.shape, 'float32')
full_samples_src = np.zeros(filled.shape, 'float32')
arg.start = ut.make_mod(arg.start, (1. / pr.fps))
ts = np.arange(arg.start, arg.start + full_dur - arg.clip_dur, step_dur)
full_ims = [None] * int(np.ceil(full_dur * pr.fps))
net = NetClf(pr, gpu=gpus[0])
for t in ut.time_est(ts):
t = ut.make_mod(t, (1. / pr.fps))
frame_start = int(t * pr.fps - arg.start * pr.fps)
ret = run(arg.vid_file,
t,
arg.clip_dur,
pr,
gpus[0],
mask=arg.mask,
arg=arg,
net=net)
if ret is None:
continue
ims = ret['ims']
for frame, im in zip(xrange(frame_start, frame_start + len(ims)), ims):
full_ims[frame] = im
def train(pr, gpus, restore=False, use_reg=True):
print 'Params:'
print pr
gpus = set_gpus(gpus)
ut.mkdir(pr.resdir)
ut.mkdir(pr.dsdir)
ut.mkdir(pr.train_dir)
config = tf.ConfigProto(allow_soft_placement=True)
if pr.inputs == ['press']:
return train_press(pr)
with tf.Graph().as_default(), tf.device(
gpus[0]), tf.Session(config=config) as sess:
global_step = tf.get_variable('global_step', [],
initializer=tf.constant_initializer(0),
trainable=False)
inputs = read_data(pr, len(gpus))
lr = pr.base_lr * pr.lr_gamma**(global_step // pr.step_size)
#opt = tf.train.MomentumOptimizer(lr, 0.9)
if pr.opt_method == 'adam':
opt = tf.train.AdamOptimizer(lr)
elif pr.opt_method == 'momentum':
opt = tf.train.MomentumOptimizer(lr, 0.9)
gpu_grads = []
for gi, gpu in enumerate(gpus):
with tf.device(gpu):
label = inputs[gi]['is_gripping']
logits = make_model(inputs[gi], pr, train=True, reuse=(gi > 0))
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=logits, labels=label)
loss = cls_loss = tf.reduce_mean(loss)
if use_reg:
reg_losses = tf.get_collection(
tf.GraphKeys.REGULARIZATION_LOSSES)
print 'Number of regularization losses:', len(reg_losses)
loss = loss + tf.add_n(reg_losses)
eq = tf.equal(tf.argmax(logits, 1), label)
acc = tf.reduce_mean(tf.cast(eq, tf.float32))
gpu_grads.append(opt.compute_gradients(loss))
#train_op = opt.minimize(loss, global_step = global_step)
grads = average_grads(gpu_grads)
train_op = opt.apply_gradients(grads, global_step=global_step)
bn_ups = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
print 'Batch norm updates:', len(bn_ups)
train_op = tf.group(train_op, *bn_ups)
sess.run(tf.global_variables_initializer())
var_list = slim.get_variables_to_restore()
exclude = [
'Adam', 'beta1_power', 'beta2_power', 'Momentum', 'global_step',
'logits', 'fc8', 'fc6_', 'fc7_', 'conv6'
]
var_list = [x for x in var_list if \
not any(name in x.name for name in exclude)]
if restore:
tf.train.Saver(var_list).restore(
sess, tf.train.latest_checkpoint(pr.train_dir))
else:
#tf.train.Saver(var_list).restore(sess, init_path)
for v in var_list:
print v.name
for base in ['im', 'depth', 'gel']:
print 'Restoring:', base
mapping = {}
for v in var_list:
print v.name
#start = '%s_resnet_v1_50/' % base
start = '%s_%s/' % (base, net_prefix())
if v.name.startswith(start):
vgg_name = v.name.replace(start, net_prefix() + '/')
vgg_name = vgg_name[:-2]
print vgg_name, '->', v.name
mapping[vgg_name] = v
if len(mapping):
tf.train.Saver(mapping).restore(sess, net_init_path())
#saver = tf.train.Saver()
tf.train.start_queue_runners(sess=sess)
summary_dir = ut.mkdir('../results/summary')
print 'tensorboard --logdir=%s' % summary_dir
sum_writer = tf.summary.FileWriter(summary_dir, sess.graph)
for i in ut.time_est(range(pr.train_iters)):
step = int(sess.run(global_step))
if (step == 10 or step % checkpoint_iters
== 0) or step == pr.train_iters - 1:
check_path = pj(ut.mkdir(pr.train_dir), 'net.tf')
print 'Saving:', check_path
vs = slim.get_model_variables()
tf.train.Saver(vs).save(sess,
check_path,
global_step=global_step)
if step > pr.train_iters:
break
merged = tf.summary.merge_all()
if step % 1 == 0:
[summary] = sess.run([merged])
sum_writer.add_summary(summary, step)
_, lr_val, cls_loss_val, loss_val, acc_val = sess.run(
[train_op, lr, cls_loss, loss, acc])
if step % 10 == 0:
print 'Iteration %d,' % step, 'lr = ', lr_val, \
'loss:', moving_avg('full_loss', loss_val), \
'cls_loss:', moving_avg('cls_loss', cls_loss_val), 'acc:', moving_avg('acc', acc_val)
sys.stdout.flush()
def write_data(out_dir,
rebalance_data=True,
train_frac=0.75,
val_frac=0.0,
n=None,
seed=0):
#def write_data(out_dir, rebalance_data = True, train_frac = 0.75, val_frac = 0.0, n = 10):
assert not os.path.exists(out_dir)
ut.mkdir(out_dir)
base_data = '../data/grasp/'
ut.sys_check('find -L %s -name "*.hdf5" > %s/all_db_files.txt' %
(base_data, out_dir))
all_db_files = map(os.path.abspath,
ut.read_lines(pj(out_dir, 'all_db_files.txt'))[:n])
all_db_files = ut.shuffled_with_seed(all_db_files, seed)
all_db_files = filter(db_ok, all_db_files)
ut.write_lines(pj(out_dir, 'db_files.txt'), all_db_files)
by_name = ut.accum_dict((name_from_file(x), x) for x in all_db_files)
names = ut.shuffled_with_seed(sorted(by_name.keys()), seed)
num_names = len(names)
num_train = int(train_frac * num_names)
num_val = int(val_frac * num_names)
i = 0
train_names = names[i:num_train]
i += num_train
val_names = names[i:i + num_val]
i += num_val
test_names = names[i:]
print num_train, num_val, len(test_names)
splits = [('train', train_names), ('val', val_names), ('test', test_names)]
print 'Number of objects in each split:'
for s, o in splits:
print s, '->', len(o)
#press_clf = press.NetClf(press_model_file, gpu = write_data_gpu)
press_clf = None #press.NetClf(press_model_file, gpu = write_data_gpu)
for dset_name, names in splits:
ut.write_lines(pj(out_dir, '%s_objects.txt' % dset_name), names)
tf_file = pj(out_dir, '%s.tf' % dset_name)
pk_file = pj(out_dir, '%s.pk' % dset_name)
full_pk_file = pj(out_dir, 'full_%s.pk' % dset_name)
if os.path.exists(tf_file):
os.remove(tf_file)
writer = tf.python_io.TFRecordWriter(tf_file)
split_db_files = ut.flatten(by_name[name] for name in names)
split_db_files = ut.shuffled_with_seed(split_db_files, dset_name)
data = []
for db_file in ut.time_est(split_db_files):
with h5py.File(db_file, 'r') as db:
#print 'keys =', db.keys()
def im(x, crop=False, compress=True):
x = ig.uncompress(x)
x = np.array(x)
if crop:
x = crop_kinect(x)
#ig.show(x)
x = ig.scale(x, (256, 256), 1)
if compress:
x = ig.compress(x)
return x
def depth(x):
x = np.array(x).astype('float32')
x = ig.scale(x, (256, 256), 1)
return x
def parse_ee(x):
names = [
'angle_of_EE_at_grasping', 'location_of_EE_at_grasping'
]
vs = [x[name].value for name in names]
ee = np.concatenate([np.array(v).flatten()
for v in vs]).astype('float32')
return ee
label_file = pj(
label_path,
db_file.split('/')[-1].replace('.hdf5', '.txt'))
if os.path.exists(label_file):
print 'Reading label from file'
is_gripping = bool(ut.read_file(label_file))
else:
is_gripping = int(np.array(db['is_gripping']))
pre, mid, _ = milestone_frames(db)
# Estimate the probability that the robot is initially gripping the object
if 0:
press_a = press_clf.predict(
im(db['/GelSightA_image'].value[mid], compress=False),
im(db['/GelSightA_image'].value[pre], compress=False))
press_b = press_clf.predict(
im(db['/GelSightB_image'].value[mid], compress=False),
im(db['/GelSightB_image'].value[pre], compress=False))
initial_press_prob = 0.5 * (press_a + press_b)
else:
initial_press_prob = np.float32(-1.)
#print initial_press_prob, ig.show(im(db['/GelSightA_image'].value[mid], compress = False))
d = dict(
gel0_pre=im(db['/GelSightA_image'].value[pre]),
gel1_pre=im(db['/GelSightB_image'].value[pre]),
gel0_post=im(db['/GelSightA_image'].value[mid]),
gel1_post=im(db['/GelSightB_image'].value[mid]),
im0_pre=im(db['/color_image_KinectA'].value[pre],
crop=True),
im0_post=im(db['/color_image_KinectA'].value[mid],
crop=True),
im1_pre=im(db['/color_image_KinectB'].value[pre],
crop=True),
im1_post=im(db['/color_image_KinectB'].value[mid],
crop=True),
depth0_pre=depth(
crop_kinect(db['/depth_image_KinectA'].value[pre])),
depth0_post=depth(
crop_kinect(db['/depth_image_KinectA'].value[mid])),
initial_press_prob=initial_press_prob,
is_gripping=int(is_gripping),
end_effector=parse_ee(db),
object_name=str(np.array(db['object_name'].value)[0]),
db_file=db_file)
data.append(d)
# for db files
ut.save(full_pk_file, data)
# rebalance data?
if rebalance_data:
by_label = [[], []]
for x in ut.shuffled_with_seed(data, 'rebalance1'):
by_label[x['is_gripping']].append(x)
n = min(map(len, by_label))
print len(data), 'before rebalance'
data = ut.shuffled_with_seed(by_label[0][:n] + by_label[1][:n],
'rebalance2')
print len(data), 'after rebalance'
writer = tf.python_io.TFRecordWriter(tf_file)
for d in data:
fbl = lambda x: tf.train.Feature(bytes_list=tf.train.BytesList(
value=[x]))
fl = lambda x: tf.train.Feature(float_list=tf.train.FloatList(
value=map(float, x.flatten())))
il = lambda x: tf.train.Feature(int64_list=tf.train.Int64List(value
=x))
feat = {
'gel0_pre': fbl(d['gel0_pre']),
'gel1_pre': fbl(d['gel1_pre']),
'gel0_post': fbl(d['gel0_post']),
'gel1_post': fbl(d['gel1_post']),
'im0_pre': fbl(d['im0_pre']),
'im0_post': fbl(d['im0_post']),
'im1_pre': fbl(d['im1_pre']),
'im1_post': fbl(d['im1_post']),
'depth0_pre': fl(d['depth0_pre']),
'depth0_post': fl(d['depth0_post']),
'end_effector': fl(d['end_effector']),
'initial_press_prob': fl(d['initial_press_prob']),
'is_gripping': il([d['is_gripping']])
}
ex = tf.train.Example(features=tf.train.Features(feature=feat))
writer.write(ex.SerializeToString())
writer.close()
ut.save(pk_file, data)
print dset_name, '->', len(data), 'examples'
