first commit

syguan96 · May 24, 2022 · 7d5a3fa · 7d5a3fa
1 parent 7139fcd
commit 7d5a3fa
Show file tree

Hide file tree

Showing 10 changed files with 111 additions and 21 deletions.
diff --git a/README.md b/README.md
@@ -3,7 +3,7 @@ Code reposity for this paper:
 **NeuroFluid: Fluid Dynamics Grounding with Particle-Driven Neural Radiance Fields.**  
 [Shanyan Guan](https://syguan96.github.io/), Huayu Deng, [Yunbo Wang](https://wyb15.github.io/)<sup>†</sup>, [Xiaokang Yang](https://scholar.google.com/citations?user=yDEavdMAAAAJ&hl=zh-CN)  
 ICML 2022   
-[[Paper]](https://github.com/syguan96/NeuroFluid)/[[Project Page]](https://syguan96.github.io/NeuroFluid/)
+[[Paper]](https://arxiv.org/pdf/2203.01762.pdf) [[Project Page]](https://syguan96.github.io/NeuroFluid/)
 
 Please cite our paper (pdf) if you find this code useful:
 ```

diff --git a/configs/bunny_warmup.yaml b/configs/bunny_warmup.yaml
@@ -0,0 +1,47 @@
+TRAIN:
+  pretained_renderer: ''
+  seed: 10
+  partial_load: False
+  N_iters: 100000
+  precrop_iters: 500
+  LR:
+    lr: 5.0e-4
+    use_scheduler: True
+    decay_epochs: 10000
+  imgW: 400
+  imgH: 400
+  scale: 1.0
+  save_interval: 5000
+  log_interval: 10
+
+TEST:
+  camera_angle_x: 0.32
+  imgW: 400
+  imgH: 400
+  scale: 1.0
+  data_path: '/data/syguan/fluid_simulation/release/data/benchmarks/honeycone/test/view_0/train/particles'
+  start_index: 0
+  end_index: 50
+  near: 9.0
+  far: 13.0
+
+# configs of renderer
+RENDERER:
+  scale: 1
+  use_mask: True
+  ray:
+    ray_chunk: 1024
+    N_importance: 128  # num of fine sample point
+    N_samples: 64    # num of coarse sample point
+  NN_search:
+    fix_radius: True
+    particle_radius: 0.025
+    search_raduis_scale: 9.0
+    N_neighbor: 20
+  encoding:
+    density: True
+    var: True
+    smoothed_pos: True
+    smoothed_dir: True
+    exclude_ray: True
+    same_smooth_factor: False
diff --git a/configs/default/end2end.yaml b/configs/default/end2end.yaml
@@ -1,19 +1,17 @@
 TRAIN:
   seed: 10
-  pretained_renderer: ''
-  pretrained_transition_model: ''
+  partial_load: False
+  pretained_renderer: 'exps/watercube/warmup/models/99999.pt'
+  pretrained_transition_model: '/data/syguan/fluid_simulation/release/ckpts/trans/transv3_694999.pt' # '/data/syguan/fluid_simulation/release/ckpts/trans/dlf.pt'
   seperate_render_transition: True # seperate set learning scheduler for renderer and transition model
   LR:
-    trans_lr: 1.0e-6
-    renderer_lr: 5.0e-4
+    trans_lr: 1.0e-5
+    renderer_lr: 3.0e-4
     use_scheduler: True
     decay_eopchs: 10000
-  pretrained:
-    trans: ''
-    renderer: ''
   N_precrop_iter: 500
   half_res: False
-  epochs: 200
+  epochs: 20000
   precrop_iters: 500
   save_interval: 500
   log_interval: 10
@@ -33,11 +31,11 @@ TEST:
 
 # configs of renderer
 RENDERER:
-  scale: 1
+  # scale: 1.0
   use_mask: True
   ray:
     ray_chunk: 1024
-    N_importance: 0  # num of fine sample point
+    N_importance: 128  # num of fine sample point
     N_samples: 64    # num of coarse sample point
   NN_search:
     fix_radius: True

diff --git a/configs/default/warmup.yaml b/configs/default/warmup.yaml
@@ -15,19 +15,19 @@ TRAIN:
   log_interval: 10
 
 TEST:
-  camera_angle_x: 0.2988457977771759
+  camera_angle_x: 0.323
   imgW: 400
   imgH: 400
   scale: 1.0
-  data_path: '/data/syguan/fluid_simulation/release/data/benchmarks/bunny/view_0/train/particles'
+  data_path: '/data/syguan/fluid_simulation/release/data/benchmarks/honeycone/test/view_0/train/particles'
   start_index: 0
   end_index: 50
   near: 9.0
   far: 13.0
 
 # configs of renderer
 RENDERER:
-  scale: 1
+  # scale: 1
   use_mask: True
   ray:
     ray_chunk: 1024

diff --git a/configs/honeycone_warmup.yaml b/configs/honeycone_warmup.yaml
@@ -19,7 +19,7 @@ TEST:
   imgW: 400
   imgH: 400
   scale: 1.0
-  data_path: '/data/syguan/fluid_simulation/release/data/benchmarks/bunny/view_0/train/particles'
+  data_path: '/data/syguan/fluid_simulation/release/data/benchmarks/honeycone/test/view_0/test/particles'
   start_index: 0
   end_index: 50
   near: 9.0

diff --git a/configs/watercube_warmup.yaml b/configs/watercube_warmup.yaml
@@ -19,7 +19,7 @@ TEST:
   imgW: 400
   imgH: 400
   scale: 1.0
-  data_path: '/data/syguan/fluid_simulation/release/data/benchmarks/bunny/view_0/train/particles'
+  data_path: '/data/syguan/fluid_simulation/release/data/benchmarks/watercube/test/view_0/test/particles'
   start_index: 0
   end_index: 50
   near: 9.0

diff --git a/configs/watersphere_warmup.yaml b/configs/watersphere_warmup.yaml
@@ -0,0 +1,47 @@
+TRAIN:
+  pretained_renderer: ''
+  seed: 10
+  partial_load: False
+  N_iters: 100000
+  precrop_iters: 500
+  LR:
+    lr: 5.0e-4
+    use_scheduler: True
+    decay_epochs: 10000
+  imgW: 400
+  imgH: 400
+  scale: 1.0
+  save_interval: 5000
+  log_interval: 10
+
+TEST:
+  camera_angle_x: 0.32
+  imgW: 400
+  imgH: 400
+  scale: 1.0
+  data_path: '/data/syguan/fluid_simulation/release/data/benchmarks/watersphere/test/view_0/test/particles'
+  start_index: 0
+  end_index: 50
+  near: 9.0
+  far: 13.0
+
+# configs of renderer
+RENDERER:
+  scale: 1
+  use_mask: True
+  ray:
+    ray_chunk: 1024
+    N_importance: 128  # num of fine sample point
+    N_samples: 64    # num of coarse sample point
+  NN_search:
+    fix_radius: True
+    particle_radius: 0.025
+    search_raduis_scale: 9.0
+    N_neighbor: 20
+  encoding:
+    density: True
+    var: True
+    smoothed_pos: True
+    smoothed_dir: True
+    exclude_ray: True
+    same_smooth_factor: False
diff --git a/train_renderer.py b/train_renderer.py
@@ -4,13 +4,11 @@
 
 import torch
 
-# from options.renderer_options_train import TrainOptions
 from trainer.trainer_renderer import Trainer
 from configs import dataset_config, warmup_training_config
 
 if __name__ == '__main__':
     torch.set_default_tensor_type('torch.cuda.FloatTensor')
-    # options = TrainOptions().parse_args()
     cfg_datasets = dataset_config()
     cfg_warmup = warmup_training_config()
 

diff --git a/trainer/trainer_e2e.py b/trainer/trainer_e2e.py
@@ -168,7 +168,7 @@ def train(self,):
         self.transition_model.train()
         self.renderer.train()
 
-        for epoch_idx in tqdm(range(self.start_step, self.options.TRAIN.epochs), total=self.options.TRAIN.epochs, desc='Iteration:'):
+        for epoch_idx in tqdm(range(self.start_step, self.options.TRAIN.epochs), total=self.options.TRAIN.epochs, desc='Epoch:'):
             self.tmp_fluid_error = FluidErrors()
             for data_idx in range(self.dataset_length):
                 data = self.dataset[data_idx]
@@ -304,7 +304,7 @@ def eval(self, step_idx):
         """
         visulize the point cloud resutls, and the image
         """
-        print('\nEval:', step_idx)
+        print('\nStep {} Eval:'.format(step_idx))
         self.eval_count += 1
         self.transition_model.eval()
         self.renderer.eval()

diff --git a/trainer/trainer_renderer.py b/trainer/trainer_renderer.py
@@ -131,7 +131,7 @@ def train_step(self, data, view_num, H, W, step_idx):
             total_loss = total_loss+rgbloss
 
             # log
-            if (step_idx+1) % self.options.log_interval == 0:
+            if (step_idx+1) % self.options.TRAIN.log_interval == 0:
                 self.summary_writer.add_scalar(f'{view_name}/rgbloss_0', rgbloss_0.item(), step_idx)
                 self.summary_writer.add_histogram(f'{view_name}/num_neighbors_0', render_ret['num_nn_0'], step_idx)
                 if N_importance>0: