def test_inference(self): """Test that inference runs and produces output of the right size.""" image1 = np.random.randn(256, 256, 3).astype('float32') image2 = np.random.randn(256, 256, 3).astype('float32') uflow = UFlow() flow = uflow.infer(image1, image2) correct_shape = np.equal(flow.shape, [256, 256, 2]).all() self.assertTrue(correct_shape)
def test_train_step(self): """Test a single training step.""" ds = tf.data.Dataset.from_tensor_slices( (tf.zeros([1, 2, 256, 256, 3], dtype=tf.float32), {})).repeat().batch(1) it = tf.compat.v1.data.make_one_shot_iterator(ds) uflow = UFlow() log_update = uflow.train(it, num_steps=1) self.assertNotEmpty(log_update)
def convert_graph(checkpoint_dir, height, width, output_row_col): """Converts uflow to a frozen model.""" print('Building the model.') tf.compat.v1.reset_default_graph() uflow = UFlow(checkpoint_dir=checkpoint_dir, channel_multiplier=FLAGS.channel_multiplier, num_levels=FLAGS.num_levels) image1 = tf.compat.v1.placeholder(tf.float32, [height, width, 3], name='first_image') image2 = tf.compat.v1.placeholder(tf.float32, [height, width, 3], name='second_image') flow = uflow.infer_no_tf_function(image1, image2, input_height=FLAGS.resize_to_height, input_width=FLAGS.resize_to_width) if output_row_col: flow = tf.identity(flow, 'flow') else: flow = tf.identity(flow[Ellipsis, ::-1], 'flow') print('Loading the checkpoint.') saver = tf.compat.v1.train.Saver() sess = tf.compat.v1.Session() sess.run(tf.compat.v1.global_variables_initializer()) sess.run(tf.compat.v1.local_variables_initializer()) # Apparently, uflow.restore() does not work here in graph mode. saver.restore(sess, tf.train.latest_checkpoint(checkpoint_dir)) # Load images. image1_np = cv2.imread( FLAGS.image1_path)[:, :, ::-1].astype('float32') / 255. image2_np = cv2.imread( FLAGS.image2_path)[:, :, ::-1].astype('float32') / 255. image1_np = cv2.resize(image1_np, (FLAGS.width, FLAGS.height)) image2_np = cv2.resize(image2_np, (FLAGS.width, FLAGS.height)) # Compute test flow and save. flow_np = sess.run(flow, feed_dict={image1: image1_np, image2: image2_np}) uflow_plotting.plot_flow(image1_np, image2_np, flow_np, 'test_result.png', checkpoint_dir) print('Freezing and optimizing the graph.') dg = tf.compat.v1.get_default_graph() frozen_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants( sess, # The session is used to retrieve the weights. dg.as_graph_def(), # The graph_def is used to retrieve the nodes. ['flow'] # The output node names are used to select the useful nodes. ) print('Writing the result to', checkpoint_dir) filename = os.path.join(checkpoint_dir, 'uflow.pb') with tf.io.gfile.GFile(filename, 'wb') as f: f.write(frozen_graph_def.SerializeToString())
def _export_frozen_graph(): # save weights tempdir = tempfile.mkdtemp() sess = tf.compat.v1.Session() with sess.as_default(): uflow = UFlow(checkpoint_dir=tempdir) image1 = tf.compat.v1.placeholder(tf.float32, [HEIGHT, WIDTH, 3], name='first_image') image2 = tf.compat.v1.placeholder(tf.float32, [HEIGHT, WIDTH, 3], name='second_image') flow = tf.identity(uflow.infer_no_tf_function(image1, image2), 'flow') sess.run(tf.compat.v1.global_variables_initializer()) image1_np = np.random.randn(HEIGHT, WIDTH, 3) image2_np = np.random.randn(HEIGHT, WIDTH, 3) flow_output = sess.run(flow, feed_dict={ image1: image1_np, image2: image2_np }) uflow.save() # freeze model freeze_model.convert_graph(tempdir, HEIGHT, WIDTH, output_row_col=True) return tempdir, image1_np, image2_np, flow_output
def create_uflow(): """Build the uflow model.""" build_selfsup_transformations = partial( uflow_augmentation.build_selfsup_transformations, crop_height=FLAGS.selfsup_crop_height, crop_width=FLAGS.selfsup_crop_width, max_shift_height=FLAGS.selfsup_max_shift, max_shift_width=FLAGS.selfsup_max_shift, resize=FLAGS.resize_selfsup) # Define learning rate schedules [none, cosine, linear, expoential]. def learning_rate_fn(): step = tf.compat.v1.train.get_or_create_global_step() effective_step = tf.maximum(step - FLAGS.lr_decay_after_num_steps + 1, 0) lr_step_ratio = tf.cast(effective_step, 'float32') / float( FLAGS.lr_decay_steps) if FLAGS.lr_decay_type == 'none' or FLAGS.lr_decay_steps <= 0: return FLAGS.gpu_learning_rate elif FLAGS.lr_decay_type == 'cosine': x = np.pi * tf.minimum(lr_step_ratio, 1.0) return FLAGS.gpu_learning_rate * (tf.cos(x) + 1.0) / 2.0 elif FLAGS.lr_decay_type == 'linear': return FLAGS.gpu_learning_rate * tf.maximum( 1.0 - lr_step_ratio, 0.0) elif FLAGS.lr_decay_type == 'exponential': return FLAGS.gpu_learning_rate * 0.5**lr_step_ratio else: raise ValueError('Unknown lr_decay_type', FLAGS.lr_decay_type) occ_weights = { 'fb_abs': FLAGS.occ_weights_fb_abs, 'forward_collision': FLAGS.occ_weights_forward_collision, 'backward_zero': FLAGS.occ_weights_backward_zero, } # Switch off loss-terms that have weights < 1e-2. occ_weights = {k: v for (k, v) in occ_weights.items() if v > 1e-2} occ_thresholds = { 'fb_abs': FLAGS.occ_thresholds_fb_abs, 'forward_collision': FLAGS.occ_thresholds_forward_collision, 'backward_zero': FLAGS.occ_thresholds_backward_zero, } occ_clip_max = { 'fb_abs': FLAGS.occ_clip_max_fb_abs, 'forward_collision': FLAGS.occ_clip_max_forward_collision, } uflow = UFlow( checkpoint_dir=FLAGS.checkpoint_dir, optimizer=FLAGS.optimizer, learning_rate=learning_rate_fn, only_forward=FLAGS.only_forward, level1_num_layers=FLAGS.level1_num_layers, level1_num_filters=FLAGS.level1_num_filters, level1_num_1x1=FLAGS.level1_num_1x1, dropout_rate=FLAGS.dropout_rate, build_selfsup_transformations=build_selfsup_transformations, fb_sigma_teacher=FLAGS.fb_sigma_teacher, fb_sigma_student=FLAGS.fb_sigma_student, train_with_supervision=FLAGS.use_supervision, train_with_gt_occlusions=FLAGS.use_gt_occlusions, smoothness_edge_weighting=FLAGS.smoothness_edge_weighting, teacher_image_version=FLAGS.teacher_image_version, stop_gradient_mask=FLAGS.stop_gradient_mask, selfsup_mask=FLAGS.selfsup_mask, normalize_before_cost_volume=FLAGS.normalize_before_cost_volume, original_layer_sizes=FLAGS.original_layer_sizes, shared_flow_decoder=FLAGS.shared_flow_decoder, channel_multiplier=FLAGS.channel_multiplier, num_levels=FLAGS.num_levels, use_cost_volume=FLAGS.use_cost_volume, use_feature_warp=FLAGS.use_feature_warp, accumulate_flow=FLAGS.accumulate_flow, occlusion_estimation=FLAGS.occlusion_estimation, occ_weights=occ_weights, occ_thresholds=occ_thresholds, occ_clip_max=occ_clip_max, smoothness_at_level=FLAGS.smoothness_at_level, ) return uflow