tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-13 19:33:46 +08:00
Code Activity

Fix bug related to distributed data parallel training (#1856)

Browse Source 
* Fix bug related to distributed data parallel training

* fix format by black

* Add nerfacto-big multi-gpu commands

* cleaned up docs

* reverted default dataset to poster

---------

Co-authored-by: Ethan Weber <eweb0124@gmail.com>
This commit is contained in:
J.YandEthan Weber authored and GitHub committed 2023-05-15 14:23:29 -07:00
1 parent f1b60d4641
commit 7851dfa9db
3 files changed
+59 -17
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/quickstart/first_nerf.md
+54 -14
View File
@@ -112,20 +112,60 @@ We also provide a train and evaluation script that allows you to do benchmarking
## Multi-GPU Training
Here we explain how to use multi-GPU training. This is the command to train the nerfacto model with 4 GPUs. We are using [PyTorch Distributed Data Parallel (DDP)](https://pytorch.org/tutorials/beginner/dist_overview.html), so we scale the learning rate with the number of GPUs used. Plotting will only be done for the first process. Note that you may want to play around with both learning rate and `<X>_num_rays_per_batch` when using DDP. Below is a simple example for how you'd run the vanilla-nerf method with either 1 or 4 GPUs. We don't see much value at the moment for running the fast methods like nerfacto with more than one GPU.
Here we explain how to use multi-GPU training. We are using [PyTorch Distributed Data Parallel (DDP)](https://pytorch.org/tutorials/beginner/dist_overview.html), so gradients are averaged over devices. If the loss scales depend on sample size (usually not the case in our implementation), we need to scale the learning rate with the number of GPUs used. Plotting will only be done for the first process. Note that you may want to play around with both learning rate and `<X>_num_rays_per_batch` when using DDP. Below is a simple example for how you'd run the `nerfacto-big` method on the aspen scene (see above to download the data), with either 1 or 2 GPUs. The `nerfacto-big` method uses a larger model size than the `nerfacto` method, so it benefits more from multi-GPU training.
First, download the aspen scene.
```python
# 1 GPU (1x LR)
export CUDA_VISIBLE_DEVICES=0
ns-train vanilla-nerf \
--machine.num-gpus 1 \
--vis wandb \
--optimizers.fields.optimizer.lr 5e-4
# 4 GPUs (4x LR)
export CUDA_VISIBLE_DEVICES=0,1,2,3
ns-train vanilla-nerf \
--machine.num-gpus 4 \
--vis wandb \
--optimizers.fields.optimizer.lr 20e-4
ns-download-data nerfstudio --capture-name=aspen
```
```python
# 1 GPU (8192 rays per GPU per batch)
export CUDA_VISIBLE_DEVICES=0
ns-train nerfacto-big --vis viewer+wandb --machine.num-gpus 1 --pipeline.datamanager.train-num-rays-per-batch 4096 --data data/nerfstudio/aspen
```
You would observe about ~70k rays / sec on NVIDIA V100.
```
Step (% Done) Train Iter (time) ETA (time) Train Rays / Sec
-----------------------------------------------------------------------------------
610 (0.61%) 115.968 ms 3 h, 12 m, 6 s 72.68 K
620 (0.62%) 115.908 ms 3 h, 11 m, 58 s 72.72 K
630 (0.63%) 115.907 ms 3 h, 11 m, 57 s 72.73 K
640 (0.64%) 115.937 ms 3 h, 11 m, 59 s 72.71 K
650 (0.65%) 115.853 ms 3 h, 11 m, 49 s 72.76 K
660 (0.66%) 115.710 ms 3 h, 11 m, 34 s 72.85 K
670 (0.67%) 115.797 ms 3 h, 11 m, 42 s 72.80 K
680 (0.68%) 115.783 ms 3 h, 11 m, 39 s 72.81 K
690 (0.69%) 115.756 ms 3 h, 11 m, 35 s 72.81 K
700 (0.70%) 115.755 ms 3 h, 11 m, 34 s 72.81 K
```
By having more GPUs in the training, you can allocate batch size to multiple GPUs and average their gradients.
```python
# 2 GPUs (4096 rays per GPU per batch, effectively 8192 rays per batch)
export CUDA_VISIBLE_DEVICES=0,1
ns-train nerfacto --vis viewer+wandb --machine.num-gpus 2 --pipeline.datamanager.train-num-rays-per-batch 4096 --data data/nerfstudio/aspen
```
You would get improved throughput (~100k rays / sec on two NVIDIA V100).
```
Step (% Done) Train Iter (time) ETA (time) Train Rays / Sec
-----------------------------------------------------------------------------------
1910 (1.91%) 79.623 ms 2 h, 10 m, 10 s 104.92 K
1920 (1.92%) 79.083 ms 2 h, 9 m, 16 s 105.49 K
1930 (1.93%) 79.092 ms 2 h, 9 m, 16 s 105.48 K
1940 (1.94%) 79.364 ms 2 h, 9 m, 42 s 105.21 K
1950 (1.95%) 79.327 ms 2 h, 9 m, 38 s 105.25 K
1960 (1.96%) 79.473 ms 2 h, 9 m, 51 s 105.09 K
1970 (1.97%) 79.334 ms 2 h, 9 m, 37 s 105.26 K
1980 (1.98%) 79.200 ms 2 h, 9 m, 23 s 105.38 K
1990 (1.99%) 79.264 ms 2 h, 9 m, 28 s 105.29 K
2000 (2.00%) 79.168 ms 2 h, 9 m, 18 s 105.40 K
```
During training, the "Train Rays / Sec" throughput represents the total number of training rays it processes per second, gradually increase the number of GPUs and observe how this throughput improves and eventually saturates.
nerfstudio/engine/trainer.py
+4 -2
View File
@@ -251,7 +251,9 @@ class Trainer:
if step > 1:
writer.put_time(
name=EventName.TRAIN_RAYS_PER_SEC,
duration=self.pipeline.datamanager.get_train_rays_per_batch() / train_t.duration,
duration=self.world_size
* self.pipeline.datamanager.get_train_rays_per_batch()
/ train_t.duration,
step=step,
avg_over_steps=True,
)
@@ -365,7 +367,7 @@ class Trainer:
train_num_rays_per_batch: int = self.pipeline.datamanager.get_train_rays_per_batch()
writer.put_time(
name=EventName.TRAIN_RAYS_PER_SEC,
duration=train_num_rays_per_batch / (train_t.duration - vis_t.duration),
duration=self.world_size * train_num_rays_per_batch / (train_t.duration - vis_t.duration),
step=step,
avg_over_steps=True,
)
nerfstudio/pipelines/base_pipeline.py
+1 -1
View File
@@ -274,7 +274,7 @@ class VanillaPipeline(Pipeline):
step: current iteration step to update sampler if using DDP (distributed)
"""
ray_bundle, batch = self.datamanager.next_train(step)
model_outputs = self.model(ray_bundle)
model_outputs = self._model(ray_bundle) # train distributed data parallel model if world_size > 1
metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
if self.config.datamanager.camera_optimizer is not None: