def download_pretrained():
# https://github.com/tensorflow/models/tree/master/slim
ut.mkdir('../results')
ut.sys_check(
'wget http://download.tensorflow.org/models/vgg_16_2016_08_28.tar.gz '
' -O ../results/vgg_16_2016_08_28.tar.gz')
ut.sys_check('cd ../results; tar -xzf vgg_16_2016_08_28.tar.gz')
def make_model(self):
with tf.device(self.default_gpu):
pr = self.pr
if self.is_training:
self.make_train_ops()
else:
self.make_test_ops(reuse=False)
self.coord = tf.train.Coordinator()
self.saver_fast = tf.train.Saver()
self.saver_slow = tf.train.Saver(max_to_keep=1000)
self.init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
self.sess.run(self.init_op)
tf.train.start_queue_runners(sess=self.sess, coord=self.coord)
print 'Initializing'
self.merged_summary = tf.summary.merge_all()
print 'Tensorboard command:'
summary_dir = ut.mkdir(pj(pr.summary_dir, ut.simple_timestamp()))
print 'tensorboard --logdir=%s' % summary_dir
self.sum_writer = tf.summary.FileWriter(summary_dir,
self.sess.graph)
if self.profile:
#self.run_meta = tf.RunMetadata()
self.profiler = tf.profiler.Profiler(self.sess.graph)
def write_data(vid_path, out_dir, train_frac=0.75):
im_dir = ut.mkdir(pj(out_dir, 'ims'))
in_data = []
meta_files = sorted(ut.glob(vid_path, 'train', '*.txt'))
print 'meta files:'
for x in meta_files:
print x
print
for meta_idx, meta_file in enumerate(meta_files):
last_prev_time = 0.
vid_file = meta_file.replace('.txt', '.mp4')
for clip_idx, ex in enumerate(ut.read_lines(meta_file)):
prev_time = last_prev_time
vid_idx = '%05d_%05d' % (meta_idx, clip_idx)
print ex
s, time = ex.split()
time = float(time)
if s == 'p':
label = 1
elif s == 'n':
label = 0
last_prev_time = time
else:
raise RuntimeError()
in_data.append((vid_file, time, label, vid_idx, im_dir, prev_time))
print 'Writing:', len(in_data), 'sequences'
meta_examples = ut.flatten(ut.parmap(extract_frames, in_data))
meta_examples = ut.shuffled_with_seed(meta_examples)
# add manu examples
db_files = sorted(
ut.sys_with_stdout('find ../data/manu-press -name "*.hdf5"').split())
db_files = ut.shuffled_with_seed(db_files)
print 'Train fraction:', train_frac
num_train = int(train_frac * len(db_files))
db_train = db_files[:num_train]
db_test = db_files[num_train:]
train_db_examples = ut.flatten(
ut.parmap(examples_from_db, [(x, im_dir) for x in db_train]))
test_db_examples = ut.flatten(
ut.parmap(examples_from_db, [(x, im_dir) for x in db_test]))
print 'Number of db train examples:', len(train_db_examples)
print 'Number of meta examples:', len(meta_examples)
train_examples = ut.shuffled_with_seed(meta_examples + train_db_examples)
ut.write_lines(pj(out_dir, 'train.csv'),
['%s,%s,%d,%s' % x for x in train_examples])
test_examples = ut.shuffled_with_seed(test_db_examples)
ut.write_lines(pj(out_dir, 'test.csv'),
['%s,%s,%d,%s' % x for x in test_examples])
def train_dir(self):
return ut.mkdir(pj(self.resdir, 'training'))
def train(path, restore=False):
config = tf.ConfigProto(allow_soft_placement=True)
with tf.Graph().as_default(), tf.device(gpu), tf.Session(
config=config) as sess:
global_step = tf.get_variable('global_step', [],
initializer=tf.constant_initializer(0),
trainable=False)
ims, ims_prev, labels = read_data(path)
#tf.summary.image('im', ims)
logits = make_model(ims, ims_prev, train=True)
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
labels=labels)
#logits = tf.Print(logits, ['logits =', logits[0, :], labels[0]])
loss = tf.reduce_mean(loss)
tf.summary.scalar('loss', loss)
eq = tf.equal(tf.argmax(logits, 1), labels)
acc = tf.reduce_mean(tf.cast(eq, tf.float32))
tf.summary.scalar('acc', acc)
lr = base_lr * gamma**(global_step // step_size)
opt = tf.train.MomentumOptimizer(lr, 0.9)
train_op = opt.minimize(loss, 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)]
train_dir = pj(path, 'training')
if restore:
tf.train.Saver().restore(sess,
tf.train.latest_checkpoint(train_dir))
else:
tf.train.Saver(var_list).restore(sess, 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)
while True:
step = int(sess.run(global_step))
if (step == 10 or step % checkpoint_iters
== 0) or step == train_iters - 1:
check_path = pj(ut.mkdir(train_dir), 'net.tf')
print 'Saving:', check_path
#saver.save(sess, check_path, global_step = global_step)
vs = slim.get_model_variables()
# print 'Variables:'
# for x in vs:
# print x.name
tf.train.Saver(vs).save(sess,
check_path,
global_step=global_step)
if step > train_iters:
break
merged = tf.summary.merge_all()
if step % 1 == 0:
[summary] = sess.run([merged])
sum_writer.add_summary(summary, step)
_, lr_val, loss_val, acc_val = sess.run([train_op, lr, loss, acc])
if step % 10 == 0:
print 'Iteration %d,' % step, 'lr = ', lr_val, 'loss:', \
moving_avg('loss', loss_val), 'acc:', moving_avg('acc', acc_val)
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'
def checkpoint_slow(self): check_path = pj(ut.mkdir(pj(self.pr.train_dir, 'slow')), 'net.tf') out = self.saver_slow.save(self.sess, check_path, global_step = self.step) print 'Checkpoint:', out
def depth_v2():
pr = base_v2().updated(
description='Depth only',
resdir=ut.mkdir('../results/grasp-params/v2/depth-v2'),
inputs=['depth'])
return pr
def make_train_model(self):
with tf.device(self.default_gpu):
pr = self.pr
# steps
self.step = tf.get_variable(
'global_step', [], trainable = False,
initializer = tf.constant_initializer(0), dtype = tf.int64)
self.lr = tf.constant(pr.base_lr)
# model
scale = pr.gamma ** tf.floor(cast_float(self.step) / float(pr.step_size))
self.lr_step = pr.base_lr * scale
#lr = tf.Print(lr, [lr, lr*1e3, scale])
opt = shift.make_opt(pr.opt_method, self.lr_step, pr)
self.inputs = read_data(pr, self.gpus)
gpu_grads, gpu_losses = {}, {}
for i, gpu in enumerate(self.gpus):
with tf.device(gpu):
reuse = (i > 0)
ims = self.inputs[i]['ims']
samples = self.inputs[i]['samples']
labels = self.inputs[i]['label']
net = make_net(ims, samples, pr, reuse = reuse, train = self.is_training)
self.loss = tfu.Loss('loss')
self.loss.add_loss(shift.slim_losses_with_prefix(None), 'reg')
self.loss.add_loss_acc(label_loss(net.logits, labels), 'label')
grads = opt.compute_gradients(self.loss.total_loss())
ut.add_dict_list(gpu_grads, self.loss.name, grads)
ut.add_dict_list(gpu_losses, self.loss.name, self.loss)
if i == 0:
self.net = net
(gs, vs) = zip(*tfu.average_grads(gpu_grads['loss']))
if pr.grad_clip is not None:
gs, _ = tf.clip_by_global_norm(gs, pr.grad_clip)
gs = [tfu.print_every(gs[0], 100, ['grad norm:', tf.global_norm(gs)])] + list(gs[1:])
gvs = zip(gs, vs)
#for g, v in zip(grads, vs):
# if g[0] is not None:
# tf.summary.scalar('%s_grad_norm' % v.name, tf.reduce_sum(g[0]**2)**0.5)
# tf.summary.scalar('%s_val_norm' % v.name, tf.reduce_sum(v**2)**0.5)
#self.train_op = opt.apply_gradients(gvs, global_step = self.step)
bn_ups = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
# self.train_op = tf.group(self.train_op, *bn_ups)
with tf.control_dependencies(bn_ups):
self.train_op = opt.apply_gradients(gvs, global_step = self.step)
self.coord = tf.train.Coordinator()
self.saver_fast = tf.train.Saver()
self.saver_slow = tf.train.Saver(max_to_keep = 1000)
#self.init_op = tf.global_variables_initializer()
if self.is_training:
self.init_op = tf.group(
tf.global_variables_initializer(),
tf.local_variables_initializer())
self.sess.run(self.init_op)
tf.train.start_queue_runners(sess = self.sess, coord = self.coord)
self.merged_summary = tf.summary.merge_all()
print 'Tensorboard command:'
summary_dir = ut.mkdir(pj(pr.summary_dir, ut.simple_timestamp()))
print 'tensorboard --logdir=%s' % summary_dir
self.sum_writer = tf.summary.FileWriter(summary_dir, self.sess.graph)
if self.profile:
self.profiler = tf.profiler.Profiler(self.sess.graph)
def make_model(self):
with tf.device(self.default_gpu):
pr = self.pr
# steps
self.step = tf.get_variable('global_step', [],
trainable=False,
initializer=tf.constant_initializer(0),
dtype=tf.int64)
self.lr = tf.constant(pr.base_lr)
# model
opt = make_opt(pr.opt_method, pr.base_lr, pr)
self.inputs = read_data(pr, self.gpus)
gpu_grads, gpu_losses = {}, {}
for i, gpu in enumerate(self.gpus):
with tf.device(gpu):
reuse = (i > 0)
with tf.device('/cpu:0'):
ims = self.inputs[i]['ims']
samples_ex = self.inputs[i]['samples']
assert pr.both_examples
assert not pr.small_augment
labels = tf.random_uniform([shape(ims, 0)],
0,
2,
dtype=tf.int64,
name='labels_sample')
samples0 = tf.where(tf.equal(labels, 1),
samples_ex[:, 1], samples_ex[:, 0])
samples1 = tf.where(tf.equal(labels, 0),
samples_ex[:, 1], samples_ex[:, 0])
labels1 = 1 - labels
net0 = make_net(ims,
samples0,
pr,
reuse=reuse,
train=self.is_training)
net1 = make_net(None,
samples1,
pr,
im_net=net0.im_net,
reuse=True,
train=self.is_training)
labels = tf.concat([labels, labels1], 0)
net = ut.Struct(
logits=tf.concat([net0.logits, net1.logits], 0),
cam=tf.concat([net0.cam, net1.cam], 0),
last_conv=tf.concat([net0.last_conv, net1.last_conv],
0))
loss = mu.Loss('loss')
loss.add_loss(slim_losses_with_prefix(None), 'reg')
loss.add_loss_acc(sigmoid_loss(net.logits, labels),
'label')
grads = opt.compute_gradients(loss.total_loss())
ut.add_dict_list(gpu_grads, loss.name, grads)
ut.add_dict_list(gpu_losses, loss.name, loss)
#self.loss = loss
if i == 0:
self.net = net
self.loss = mu.merge_losses(gpu_losses['loss'])
for name, val in zip(self.loss.get_loss_names(),
self.loss.get_losses()):
tf.summary.scalar(name, val)
if not self.is_training:
#pr_test = pr.copy()
pr_test = self.pr_test.copy()
pr_test.augment_ims = False
print 'pr_test ='
print pr_test
self.test_ims, self.test_samples, self.test_ytids = mu.on_cpu(
lambda: shift_dset.make_db_reader(
pr_test.test_list,
pr_test,
pr.test_batch, ['im', 'samples', 'ytid'],
one_pass=True))
if pr_test.do_shift:
self.test_labels = tf.random_uniform(
[shape(self.test_ims, 0)], 0, 2, dtype=tf.int64)
self.test_samples = tf.where(tf.equal(self.test_labels, 1),
self.test_samples[:, 1],
self.test_samples[:, 0])
else:
self.test_labels = tf.ones(shape(self.test_ims, 0),
dtype=tf.int64)
#self.test_samples = tf.where(tf.equal(self.test_labels, 1), self.test_samples[:, 1], self.test_samples[:, 0])
print 'sample shape:', shape(self.test_samples)
self.test_net = make_net(self.test_ims,
self.test_samples,
pr_test,
reuse=True,
train=self.is_training)
(gs, vs) = zip(*mu.average_grads(gpu_grads['loss']))
if pr.grad_clip is not None:
gs, _ = tf.clip_by_global_norm(gs, pr.grad_clip)
gs = [
mu.print_every(gs[0], 100,
['grad norm:', tf.global_norm(gs)])
] + list(gs[1:])
gvs = zip(gs, vs)
bn_ups = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
if pr.multipass:
ops = [
opt.apply_gradients(gvs, global_step=self.step)
for i in xrange(pr.multipass_count)
]
def op_helper(count=[0]):
op = ops[count[0] % len(ops)]
count[0] += 1
return op
self.train_op = op_helper
else:
op = tf.group(opt.apply_gradients(gvs, global_step=self.step),
*bn_ups)
self.train_op = lambda: op
self.coord = tf.train.Coordinator()
self.saver = tf.train.Saver()
self.init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
self.sess.run(self.init_op)
tf.train.start_queue_runners(sess=self.sess, coord=self.coord)
self.merged_summary = tf.summary.merge_all()
print 'Tensorboard command:'
summary_dir = ut.mkdir(pj(pr.summary_dir, ut.simple_timestamp()))
print 'tensorboard --logdir=%s' % summary_dir
self.sum_writer = tf.summary.FileWriter(summary_dir,
self.sess.graph)
def write_data(out_dir, train_frac=0.75, val_frac=0.05):
ut.mkdir(out_dir)
base_data = '../data/grasp/'
ut.sys_check('find %s -name "*.hdf5" > %s/db_files.txt' %
(base_data, out_dir))
all_db_files = ut.read_lines(pj(out_dir, 'db_files.txt'))
all_db_files = ut.shuffled_with_seed(all_db_files)
name_from_file = lambda x: '_'.join(x.split('/')[-1].split('_')[2:])
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()))
num_names = len(all_db_files)
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:]
for dset_name, names in [('train', train_names), ('val', val_names),
('test', test_names)]:
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)
if os.path.exists(tf_file):
os.remove(tf_file)
writer = tf.python_io.TFRecordWriter(tf_file)
data = []
for name in names:
for db_file in by_name[name]:
with h5py.File(db_file, 'r') as db:
def im(x):
x = np.array(x)
x = ig.scale(x, (256, 256), 1)
return ig.compress(x)
if 'is_gripping' in db:
label = int(np.array(db['is_gripping']))
elif 'Is gripping?' in db:
label = int(np.array(db['Is gripping?']))
else:
print 'Skipping: %s. Missing is_gripping' % db_file
print 'Keys:', ' '.join(db.keys())
continue
data.append({
'gel0_pre':
im(db['GelSightA_image_pre_gripping']),
'gel1_pre':
im(db['GelSightB_image_pre_gripping']),
'gel0_post':
im(db['GelSightA_image_post_gripping']),
'gel1_post':
im(db['GelSightB_image_post_gripping']),
'is_gripping':
label
})
fbl = lambda x: tf.train.Feature(bytes_list=tf.train.
BytesList(value=[x]))
feat = {
'gel0_pre':
fbl(im(db['GelSightA_image_pre_gripping'])),
'gel1_pre':
fbl(im(db['GelSightB_image_pre_gripping'])),
'gel0_post':
fbl(im(db['GelSightA_image_post_gripping'])),
'gel1_post':
fbl(im(db['GelSightB_image_post_gripping'])),
'is_gripping':
tf.train.Feature(int64_list=tf.train.Int64List(
value=[label]))
}
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'
def shift_lowfps(num_gpus=1, shift_dur=4.2):
total_dur = 10.1
fps = 29.97
frame_dur = 1. / fps
samp_sr = 21000.
spec_sr = 100.
pr = Params(
subsample_frames=4,
train_iters=100000,
opt_method='momentum',
base_lr=1e-2,
full_model=True,
grad_clip=5.,
skip_notfound=False,
augment_ims=True,
cam=False,
batch_size=int(15 * num_gpus),
test_batch=10,
shift_dur=shift_dur,
multipass=False,
both_examples=True,
small_augment=False,
resdir=ab('/data/scratch/owens/shift/shift-lowfps'),
init_path=None,
weight_decay=1e-5,
train_list='/data/ssd1/owens/audioset-vid-v21/small_train.txt',
test_list='/data/ssd1/owens/audioset-vid-v21/small_train.txt',
num_dbs=None,
im_type='jpeg',
input_type='samples',
full_im_dim=256,
full_flow_dim=256,
crop_im_dim=224,
sf_pad=int(0.5 * 2**4 * 4),
use_flow=False,
#renorm = True,
renorm=True,
checkpoint_iters=1000,
dset_seed=None,
samp_sr=samp_sr,
spec_sr=spec_sr,
fps=fps,
#neg_frame_buf = -50,
max_intersection=30 * 2,
specgram_sr=spec_sr,
num_mel=64,
batch_norm=True,
show_videos=False,
check_iters=1000,
decompress_flow=True,
print_iters=10,
total_frames=int(total_dur * fps),
sampled_frames=int(shift_dur * fps),
full_specgram_samples=int(total_dur * spec_sr),
full_samples_len=int(total_dur * samp_sr),
sfs_per_frame=spec_sr * frame_dur,
samples_per_frame=samp_sr * frame_dur,
frame_sample_delta=int(total_dur * fps) / 2,
fix_frame=False,
use_3d=True,
augment=False,
dilate=False,
do_shift=True,
variable_frame_count=False,
momentum_rate=0.9,
use_sound=True,
bn_last=True,
summary_iters=10,
im_split=True,
num_splits=2,
augment_audio=False,
multi_shift=False,
)
pr.vis_dir = ut.mkdir(pj(pr.resdir, 'vis'))
return pr
def summary_dir(self):
return ut.mkdir(pj(self.resdir, 'summary'))
def checkpoint_fast(self): check_path = pj(ut.mkdir(self.pr.train_dir), 'net.tf') out = self.saver_fast.save(self.sess, check_path, global_step = self.step) print 'Checkpoint:', out
full_samples_src[inds[ok]] = samples_src[ok]
filled[inds] = True
full_samples_fg = np.clip(full_samples_fg, -1., 1.)
full_samples_bg = np.clip(full_samples_bg, -1., 1.)
full_samples_src = np.clip(full_samples_src, -1., 1.)
full_ims = [x for x in full_ims if x is not None]
table = [['start =', arg.start], 'fg:',
imtable.Video(full_ims, pr.fps, Sound(full_samples_fg,
pr.samp_sr)), 'bg:',
imtable.Video(full_ims, pr.fps, Sound(full_samples_bg,
pr.samp_sr)), 'src:',
imtable.Video(full_ims, pr.fps, Sound(full_samples_src,
pr.samp_sr))]
if arg.out is not None:
ut.mkdir(arg.out)
vid_s = arg.vid_file.split('/')[-1].split('.mp4')[0]
mask_s = '' if arg.mask is None else '_%s' % arg.mask
cam_s = '' if not arg.cam else '_cam'
suffix_s = '' if arg.suffix == '' else '_%s' % arg.suffix
name = '%s%s%s_%s' % (suffix_s, mask_s, cam_s, vid_s)
def snd(x):
x = Sound(x, pr.samp_sr)
x.samples = np.clip(x.samples, -1., 1.)
return x
print 'Writing to:', arg.out
ut.save(pj(arg.out, 'ret%s.pk' % name), ret)
ut.make_video(full_ims, pr.fps, pj(arg.out, 'fg%s.mp4' % name),
snd(full_samples_fg))
def gel_im_depth_v2():
pr = base_v2().updated(
description='GelSight + images + depth',
resdir=ut.mkdir('../results/grasp-params/v2/gel-im-depth-v2'),
inputs=['gel', 'im', 'depth'])
return pr
