code change

.gitattributes

@@ -1,17 +0,0 @@
-*.bin.* filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tar.gz filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

configs/diet_nerf_tpu_vm_4shot.yaml

@@ -8,8 +8,9 @@ white_bkgd: true
 batch_size: 1024
 randomized: true
 max_steps: 200000
+stop_sc_loss: 160000
 print_every: 100
-render_every: 500
+render_every: 1000
 save_every: 5000
 use_semantic_loss: true
 clip_model_name: openai/clip-vit-base-patch32

configs/diet_nerf_tpu_vm_few_shot.yaml

@@ -8,8 +8,9 @@ white_bkgd: true
 batch_size: 1024
 randomized: true
 max_steps: 200000
+stop_sc_loss: 160000
 print_every: 100
-render_every: 500
+render_every: 1000
 save_every: 5000
 use_semantic_loss: true
 clip_model_name: openai/clip-vit-base-patch32

configs/diet_nerf_tpu_vm_test.yaml

@@ -2,12 +2,13 @@ dataset: blender
 batching: single_image
 factor: 0
 num_coarse_samples: 64
-num_fine_samples: 64
+num_fine_samples: 128
 use_viewdirs: true
 white_bkgd: true
-batch_size: 1026
+batch_size: 1024
 randomized: true
 max_steps: 200000
+stop_sc_loss: 160000
 print_every: 100
 render_every: 1000
 save_every: 5000

eval.py

@@ -112,30 +112,39 @@ def main(unused_argv):
         summary_writer = tensorboard.SummaryWriter(
             path.join(FLAGS.train_dir, "eval"))
 
-    def generate_spinning_gif(radius, phi, gif_fn, frame_n):
+    def generate_spinning_gif(radius, phi, output_dir, frame_n):
         _rng = random.PRNGKey(0)
         partial_render_fn = functools.partial(render_pfn, state.optimizer.target)
         gif_images = []
+        gif_images2 = []
         for theta in tqdm(np.linspace(-math.pi, math.pi, frame_n)):
             camtoworld = np.array(clip_utils.pose_spherical(radius, theta, phi))
             rays = dataset.camtoworld_matrix_to_rays(camtoworld, downsample=4)
             _rng, key0, key1 = random.split(_rng, 3)
-            color, _, _ = utils.render_image(partial_render_fn, rays,
+            color, disp, _ = utils.render_image(partial_render_fn, rays,
                                              _rng, False, chunk=4096)
             image = predict_to_image(color)
+            image2 = predict_to_image(disp[Ellipsis, 0])
             gif_images.append(image)
+            gif_images2.append(image2)
+
+        gif_fn = os.path.join(output_dir, 'rgb_spinning.gif')
+        gif_fn2 = os.path.join(output_dir, 'disp_spinning.gif')
         gif_images[0].save(gif_fn, save_all=True,
                            append_images=gif_images,
                            duration=100, loop=0)
-        return gif_images
+        gif_images2[0].save(gif_fn2, save_all=True,
+                           append_images=gif_images2,
+                           duration=100, loop=0)
+
+        #return gif_images, gif_images2
 
     if FLAGS.generate_gif_only:
         print('generate GIF file only')
         _radius = 4.
         _phi = (30 * math.pi) / 180
-        _gif_fn = os.path.join(out_dir, 'spinning.gif')
-        generate_spinning_gif(_radius, _phi, _gif_fn, frame_n=30)
-        print(f'GIF file for spinning views written: {_gif_fn}')
+        generate_spinning_gif(_radius, _phi, out_dir, frame_n=30)
+        print('GIF file for spinning views written)')
         return
     else:
         print('generate GIF file AND evaluate model performance')
@@ -149,6 +158,7 @@ def main(unused_argv):
             utils.makedirs(out_dir)
         psnr_values = []
         ssim_values = []
+
         #lpips_values = []
         if not FLAGS.eval_once:
             showcase_index = np.random.randint(0, dataset.size)
@@ -225,9 +235,8 @@ def main(unused_argv):
             if not is_gif_written:
                 _radius = 4.
                 _phi = (30 * math.pi) / 180
-                _gif_fn = os.path.join(out_dir, 'spinning.gif')
-                generate_spinning_gif(_radius, _phi, _gif_fn, frame_n=30)
-                print(f'GIF file for spinning views written: {_gif_fn}')
+                generate_spinning_gif(_radius, _phi, out_dir, frame_n=30)
+                print(f'GIF file for spinning views written')
                 is_gif_written = True
 
         if FLAGS.eval_once:

nerf/clip_utils.py

@@ -15,50 +15,44 @@ FLAGS = flags.FLAGS
 
 @partial(jax.jit, static_argnums=[0])
 def semantic_loss(clip_model, src_image, target_embedding): 
-    #c_image = utils.unshard(src_image[0])
-    f_image = utils.unshard(src_image[-1])
-
-    w = int(math.sqrt(src_image[-1].size//3))
-    #c_image = c_image.reshape([w, w, 3])
+    c_image = utils.unshard(src_image[0])
+    f_image = utils.unshard(src_image[1])
+    w = int(math.sqrt(f_image.shape[0]))
+    c_image = c_image.reshape([w, w, 3])
     f_image = f_image.reshape([w, w, 3])
- 
-    src_embedding = clip_model.get_image_features(pixel_values=preprocess_for_CLIP(jnp.expand_dims(f_image,0).transpose(0, 3, 1, 2)))
-    #src_embedding = clip_model.get_image_features(pixel_values=preprocess_for_CLIP(jnp.stack([c_image, f_image]).transpose(0, 3, 1, 2)))
+     
+    src_embedding = clip_model.get_image_features(pixel_values=preprocess_for_CLIP(jnp.stack([c_image,f_image],0).transpose(0, 3, 1, 2)))
     src_embedding /= jnp.linalg.norm(src_embedding, axis=-1, keepdims=True)
-    sc_loss = 1 - jnp.sum(src_embedding * target_embedding)
+    sc_loss = 2 - jnp.sum(src_embedding * target_embedding)
     return sc_loss, f_image
 
 def semantic_step_multi(render_pfn, clip_model, rng, state, batch, lr):
-    random_rays = jax.tree_map(lambda x: utils.shard(x).astype(jnp.float16), batch["random_rays"])
-    target_embedding = batch["embedding"].astype(jnp.float16)
+    random_rays = batch["random_rays"]
+    target_embedding = batch["embedding"]
     rng, key_0, key_1 = random.split(rng,3)
-
+    
     def loss_fn(variables):
-        src_image = render_pfn(variables, key_0, key_1, random_rays)
-        sc_loss, f_image = semantic_loss(clip_model, src_image, target_embedding)
+        images = render_pfn(variables, key_0, key_1, random_rays)
+        sc_loss, f_image = semantic_loss(clip_model, images, target_embedding)
         return sc_loss * FLAGS.sc_loss_mult, f_image
     (sc_loss, src_image), grad = jax.value_and_grad(loss_fn, has_aux = True)(jax.device_get(jax.tree_map(lambda x:x[0], state)).optimizer.target)
     return sc_loss, grad, src_image
 
 @partial(jax.jit, static_argnums=[0, 1])
 def semantic_step_single(model, clip_model, rng, state, batch, lr):
-    batch = jax.tree_map(lambda x: x.astype(jnp.float16), batch)
-    # the batch is without shard
-    random_rays = batch["random_rays"]
+    random_rays = jax.tree_map(lambda x: x.reshape(-1,3), batch["random_rays"])
+    target_embedding = batch["embedding"]
     rng, key_0, key_1 = random.split(rng,3)
 
     def semantic_loss(variables):
         c_image, f_image = model.apply(variables, key_0, key_1, random_rays, False, rgb_only = True)
-        # reshape flat pixel to an image (assume 3 channels & square shape)
         w = int(math.sqrt(f_image.shape[0]))
-        # c_image = c_image.reshape([w, w, 3])
+        c_image = c_image.reshape([w, w, 3])
         f_image = f_image.reshape([w, w, 3])
 
-        src_embedding = clip_model.get_image_features(pixel_values=preprocess_for_CLIP(jnp.expand_dims(f_image,0).transpose(0, 3, 1, 2)))
-        # src_embedding = clip_model.get_image_features(pixel_values=preprocess_for_CLIP(jnp.stack([c_image, f_image]).transpose(0, 3, 1, 2)))
+        src_embedding = clip_model.get_image_features(pixel_values=preprocess_for_CLIP(jnp.stack([c_image,f_image],0).transpose(0, 3, 1, 2)))
         src_embedding /= jnp.linalg.norm(src_embedding, axis=-1, keepdims=True)
-        target_embedding = batch["embedding"]
-        sc_loss = 0.5 * jnp.sum((src_embedding - target_embedding)**2)
+        sc_loss = 2 - jnp.sum(src_embedding * target_embedding)
         return sc_loss * FLAGS.sc_loss_mult, f_image
     (sc_loss, src_image), grad = jax.value_and_grad(semantic_loss, has_aux = True)(jax.device_get(jax.tree_map(lambda x:x[0], state)).optimizer.target)
     return sc_loss, grad, src_image

nerf/datasets.py

@@ -236,6 +236,7 @@ class Blender(Dataset):
         camera_angle_x = float(meta["camera_angle_x"])
         self.focal = .5 * self.w / np.tan(.5 * camera_angle_x)
         self.n_examples = self.images.shape[0]
+        self.dtype = flags.clip_output_dtype
 
         if flags.use_semantic_loss and clip_model is not None:
             embs = []
@@ -258,8 +259,8 @@ class Blender(Dataset):
 
         frames = np.arange(len(meta["frames"]))
         if few_shot > 0 and split == 'train':
-            np.random.seed(0)
-            np.random.shuffle(frames)
+            # np.random.seed(0)
+            # np.random.shuffle(frames)
             frames = frames[:few_shot]
 
         # if split == 'train':
@@ -308,16 +309,21 @@ class Blender(Dataset):
         src_seed = int(time.time())
         src_rng = jax.random.PRNGKey(src_seed)
         src_camtoworld = np.array(clip_utils.random_pose(src_rng, (self.near, self.far)))
-        random_rays = self.camtoworld_matrix_to_rays(src_camtoworld, downsample = 4)
-        cx = np.random.randint(80, 120)
-        cy = np.random.randint(80, 120)
-        d = 70
-        random_rays = jax.tree_map(lambda x: x[cy-d:cy+d,cx-d:cx+d], random_rays)
+
+        cx = np.random.randint(320, 480)
+        cy = np.random.randint(320, 480)
+        d = 140
+        
+        random_rays = self.camtoworld_matrix_to_rays(src_camtoworld, downsample = 1)
+        random_rays = jax.tree_map(lambda x: x[cy-d:cy+d:4,cx-d:cx+d:4], random_rays)
+
         w = random_rays[0].shape[0] - random_rays[0].shape[0]%jax.local_device_count()
         random_rays = jax.tree_map(lambda x: x[:w,:w].reshape(-1,3), random_rays)
-        batch_dict["random_rays"] = random_rays
+        batch_dict["random_rays"] = utils.shard(random_rays)
+        if self.dtype == 'float16':
+            batch_dict = jax.tree_map(lambda x: x.astype(np.float16), batch_dict)
         return batch_dict
-
+    
 class LLFF(Dataset):
     """LLFF Dataset."""
 

nerf/models.py

@@ -136,7 +136,7 @@ class NerfModel(nn.Module):
             (comp_rgb, disp, acc),
         ]
 
-        if self.num_fine_samples > 0 and not(rgb_only):
+        if self.num_fine_samples > 0:
             z_vals_mid = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
             key, rng_1 = random.split(rng_1)
 
@@ -191,8 +191,7 @@ class NerfModel(nn.Module):
             )
             ret.append((comp_rgb, disp, acc))
         if rgb_only:
-            #return [ret[0][0], ret[1][0]]
-            return [None, ret[0][0]]
+            return [ret[0][0], ret[1][0]]
         return ret
 
 def construct_nerf(key, example_batch, args):

nerf/utils.py

@@ -66,11 +66,11 @@ def define_flags():
     flags.DEFINE_bool("use_semantic_loss", True,
                       "whether use semantic loss or not")
     flags.DEFINE_string("clip_model_name", "openai/clip-vit-base-patch32", "model type for CLIP")
-    flags.DEFINE_string("clip_output_dtype", "float32",
+    flags.DEFINE_string("clip_output_dtype", "float16",
                         "float32/ float16 (float16 for memory saving)")
     flags.DEFINE_integer("sc_loss_every", 16,
                          "no. of steps to take before performing semantic loss evaluation")
-    flags.DEFINE_float("sc_loss_mult", 1e-3,
+    flags.DEFINE_float("sc_loss_mult", 1e-2,
                        "weighting for semantic loss from CLIP")
 
     # Dataset Flags
@@ -166,6 +166,8 @@ def define_flags():
 
     flags.DEFINE_integer("max_steps", 1000000,
                          "the number of optimization steps.")
+    flags.DEFINE_integer("stop_sc_loss", 1000000,
+                         "the number of sc_loss optimization steps")
     flags.DEFINE_integer("save_every", 10000,
                          "the number of steps to save a checkpoint.")
     flags.DEFINE_integer("print_every", 100,

train.py

@@ -50,7 +50,6 @@ print(f"detected device: {jax.local_devices()}")
 
 
 def train_step(model, clip_model, rng, state, batch, lr, step, K,):
-    # TODO make clip_grad input enable
     """One optimization step.
 
     Args:
@@ -102,7 +101,6 @@ def train_step(model, clip_model, rng, state, batch, lr, step, K,):
 
     (_, stats), grad = (
         jax.value_and_grad(loss_fn, has_aux=True)(state.optimizer.target))
-    #grad = jax.lax.pmean(grad, axis_name="batch")
     stats = jax.lax.pmean(stats, axis_name="batch")
     
     # Clip the gradient by value.
@@ -238,26 +236,16 @@ def main(unused_argv):
 
         grad, stats, keys = train_pstep(keys, state, batch, lr, step, FLAGS.sc_loss_every)
         
-        if step%FLAGS.sc_loss_every == 0 and FLAGS.use_semantic_loss:
+        if step%FLAGS.sc_loss_every == 0 and FLAGS.use_semantic_loss and step < FLAGS.stop_sc_loss:
             sc_batch = dataset.get_clip_data()
             if jax.local_device_count() > 1:
                 sc_loss, sc_grad, sc_image = clip_utils.semantic_step_multi(render_pfn_, clip_model, keys[0], state, sc_batch, lr)
             else:
                 sc_loss, sc_grad, sc_image = clip_utils.semantic_step_single(model, clip_model, keys[0], state, sc_batch, lr)
 
-            if jax.host_id() == 0 and step%FLAGS.print_every:
-                for mlp_k, mlp in grad['params'].items():
-                    for layer_k, layer_g in mlp.items():
-                        summary_writer.scalar("%s/%s/kernel_grad"%(mlp_k, layer_k), jnp.linalg.norm(jnp.mean(layer_g['kernel'],0)), step)
-                for mlp_k, mlp in sc_grad['params'].items():
-                    for layer_k, layer_g in mlp.items():
-                        summary_writer.scalar("%s/%s/kernel_sc_grad"%(mlp_k, layer_k), jnp.linalg.norm(layer_g['kernel']), step)
-  
             leaves, treedef = jax.tree_flatten(grad)
             sc_leaves, _ = jax.tree_flatten(sc_grad)
             grad = treedef.unflatten(g+jnp.expand_dims(sc_g,0) for g, sc_g in zip(leaves, sc_leaves))
-            
-
         
         state = update_pstep(state, grad, lr)
 
@@ -276,24 +264,26 @@ def main(unused_argv):
                 summary_writer.scalar("psnr/train", stats.psnr[0], step)
                 summary_writer.scalar("train_coarse/loss", stats.loss_c[0], step)
                 summary_writer.scalar("train_coarse/psnr", stats.psnr_c[0], step)
-                summary_writer.scalar("weight_l2", stats.weight_l2[0], step)
+
                 avg_loss = np.mean(np.concatenate([s.loss for s in stats_trace]))
                 avg_psnr = np.mean(np.concatenate([s.psnr for s in stats_trace]))
                 stats_trace = []
                 summary_writer.scalar("train_avg/loss", avg_loss, step)
                 summary_writer.scalar("train_avg/psnr", avg_psnr, step)
-                summary_writer.scalar("learning_rate", lr, step)
+
                 steps_per_sec = FLAGS.print_every / (time.time() - t_loop_start)
                 reset_timer = True
                 rays_per_sec = FLAGS.batch_size * steps_per_sec
-                summary_writer.scalar("train_steps_per_sec", steps_per_sec, step)
-                summary_writer.scalar("train_rays_per_sec", rays_per_sec, step)
+                summary_writer.scalar("stats/weight_l2", stats.weight_l2[0], step)
+                summary_writer.scalar("stats/learning_rate", lr, step)
+                summary_writer.scalar("iter_speed/train_steps_per_sec", steps_per_sec, step)
+                summary_writer.scalar("iter_speed/train_rays_per_sec", rays_per_sec, step)
                 precision = int(np.ceil(np.log10(FLAGS.max_steps))) + 1
                 print(("{:" + "{:d}".format(precision) + "d}").format(step) +
                       f"/{FLAGS.max_steps:d}: " + f"i_loss={stats.loss[0]:0.4f}, " +
                       f"avg_loss={avg_loss:0.4f}, " +
                       f"weight_l2={stats.weight_l2[0]:0.2e}, " +
-                    #   f"sc_loss={sc_loss:0.4f}, " +
+                      f"sc_loss={sc_loss:0.4f}, " +
                       f"lr={lr:0.2e}, {rays_per_sec:0.0f} rays/sec")
             if step % FLAGS.save_every == 0:
                 state_to_save = jax.device_get(jax.tree_map(lambda x: x[0], state))
@@ -324,12 +314,10 @@ def main(unused_argv):
                 eval_time = time.time() - t_eval_start
                 num_rays = jnp.prod(jnp.array(test_case["rays"].directions.shape[:-1]))
                 rays_per_sec = num_rays / eval_time
-                summary_writer.scalar("test_rays_per_sec", rays_per_sec, step)
+                summary_writer.scalar("iter_speed/test_rays_per_sec", rays_per_sec, step)
                 print(f"Eval {step}: {eval_time:0.3f}s., {rays_per_sec:0.0f} rays/sec")
                 summary_writer.scalar("psnr/test", psnr, step)
-                summary_writer.scalar("test_psnr", psnr, step)
                 summary_writer.scalar("ssim/ssim", ssim, step)
-                summary_writer.scalar("test_ssim", ssim, step)
                 if sc_image is not None:
                     summary_writer .image("random_ray_image", sc_image, step)
                 summary_writer.image("test_pred_color", pred_color, step)