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