Merge pull request #4 from SKT-r/feat/make_set_ray_batich_size

set_ray_batch_size

Author: SakodaShintaro

localization/nerf_based_localizer/training_tool/config/train_config.yaml

@@ -5,6 +5,7 @@ train:
   end_iter: 20000
   report_freq: 200
   vis_freq: 2500
+  ray_batch_size: 4096
   save_freq: 20000
   validate_freq: 100000
   learning_rate: 1e-2

localization/nerf_based_localizer/training_tool/src/train_manager.cpp

@@ -52,6 +52,7 @@ TrainManager::TrainManager(const std::string & train_result_dir, const std::stri
   pts_batch_size_ = (int)train_config["pts_batch_size"];
   end_iter_ = (int)train_config["end_iter"];
   vis_freq_ = (int)train_config["vis_freq"];
+  ray_batch_size_ = (int)train_config["ray_batch_size"];
   report_freq_ = (int)train_config["report_freq"];
   save_freq_ = (int)train_config["save_freq"];
   learning_rate_ = (float)train_config["learning_rate"];
@@ -132,9 +133,9 @@ void TrainManager::train()
       const int idx = 0;
       const int H = dataset_->height;
       const int W = dataset_->width;
-      const int ray_batch_size = 8192;
+
       auto [pred_colors, pred_depths] = renderer_->render_image(
-        dataset_->poses[idx], dataset_->intrinsics[idx], H, W, ray_batch_size);
+        dataset_->poses[idx], dataset_->intrinsics[idx], H, W, ray_batch_size_);
 
       Tensor image = dataset_->images[idx].reshape({H, W, 3}).to(torch::kCPU);
       pred_colors = pred_colors.to(torch::kCPU);

localization/nerf_based_localizer/training_tool/src/train_manager.hpp

@@ -47,7 +47,7 @@ class TrainManager
   unsigned iter_step_ = 0;
   unsigned end_iter_;
   unsigned report_freq_, vis_freq_, save_freq_;
-  unsigned pts_batch_size_;
+  unsigned pts_batch_size_, ray_batch_size_;
 
   int var_loss_start_, var_loss_end_;
   float learning_rate_, learning_rate_alpha_, learning_rate_warm_up_end_iter_;