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

Merge pull request #45 from ethanweber/ethan/apr29_cpu_and_multigpu

Browse Source 
ethan/apr29_cpu_and_multigpu
This commit is contained in:
Ethan Weber authored and GitHub committed 2022-04-29 15:32:10 -07:00
2 parents 9004584026 023e6a3f92
commit 1f45b3e978
7 files changed
+74 -38
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.vscode/launch.json
Vendored
+9 -2
View File
@@ -12,13 +12,20 @@
"justMyCode": true
},
{
"name": "Python: run_train",
"name": "Python: run_train (gpu)",
"type": "python",
"request": "launch",
"program": "scripts/run_train.py",
"console": "integratedTerminal"
},
{
"name": "Python: run_train (cpu)",
"type": "python",
"request": "launch",
"program": "scripts/run_train.py",
"console": "integratedTerminal",
"args": [
"machine_config.num_gpus=1"
"machine_config.num_gpus=0"
]
}
]
configs/data/default_data.yaml
+3
View File
@@ -4,6 +4,9 @@ dataset:
dataset_type: blender
downscale_factor: 4
# validation options
validation_image_indices: [0, 10] # Which images to evaluate on.
# dataloader options
dataloader:
num_workers: 4
configs/default.yaml
+1 -1
View File
@@ -20,4 +20,4 @@ machine_config:
num_gpus: 1
num_machines: 1
machine_rank: 0
dist_url: auto
dist_url: auto
mattport/nerf/trainer.py
+31 -14
View File
@@ -29,7 +29,15 @@ logging.getLogger("PIL").setLevel(logging.WARNING)
class Trainer:
"""Training class"""
def __init__(self, config: DictConfig, local_rank: int = 0, world_size: int = 1):
def __init__(self, config: DictConfig, local_rank: int = 0, world_size: int = 1, cpu=False):
"""_summary_
Args:
config (DictConfig): _description_
local_rank (int, optional): _description_. Defaults to 0.
world_size (int, optional): _description_. Defaults to 1.
cpu (bool, optional): Whether or not to use the CPU.
"""
self.config = config
self.local_rank = local_rank
self.world_size = world_size
@@ -43,13 +51,13 @@ class Trainer:
self.optimizers = None
self.start_step = 0
# logging variables
self.is_main_thread = local_rank % world_size == 0
self.is_main_thread = world_size != 0 and local_rank % world_size == 0
writer = getattr(mattport.utils.writer, self.config.logging.writer.type)
self.writer = writer(local_rank, world_size, save_dir=self.config.logging.writer.save_dir)
self.writer = writer(self.is_main_thread, save_dir=self.config.logging.writer.save_dir)
self.stats = StatsTracker(config, self.is_main_thread)
if not profiler.PROFILER and self.config.logging.enable_profiler:
profiler.PROFILER = profiler.Profiler(config, self.is_main_thread)
self.dry_run = self.config.get("dry_run", False)
self.device = f"cuda:{self.local_rank}" if not cpu else "cpu"
@profiler.time_function
def setup(self):
@@ -58,6 +66,12 @@ class Trainer:
self.setup_datasets(dataset_inputs_dict)
self.setup_graph(dataset_inputs_dict["train"])
def collate_fn(self, batch_list):
"""TODO(ethan): I need to replace this.
I'm only using this for multiprocess pickle issues for now.
"""
return collate_batch(batch_list, self.config.data.dataloader.num_rays_per_batch, keep_full_image=False)
@profiler.time_function
def setup_datasets(self, dataset_inputs_dict: Dict[str, DatasetInputs]):
"""_summary_"""
@@ -65,11 +79,10 @@ class Trainer:
image_filenames=dataset_inputs_dict["train"].image_filenames,
downscale_factor=dataset_inputs_dict["train"].downscale_factor,
)
self.train_dataset = CollateIterDataset(
self.train_image_dataset,
collate_fn=lambda batch_list: collate_batch(
batch_list, self.config.data.dataloader.num_rays_per_batch, keep_full_image=False
),
collate_fn=self.collate_fn,
num_samples_to_collate=self.config.data.dataloader.num_images_to_sample_from,
num_times_to_repeat=self.config.data.dataloader.num_times_to_repeat_images,
)
@@ -96,7 +109,7 @@ class Trainer:
self.config.graph.network,
intrinsics=dataset_inputs.intrinsics,
camera_to_world=dataset_inputs.camera_to_world,
).to(f"cuda:{self.local_rank}")
).to(self.device)
self.setup_optimizers()
if self.config.graph.resume_train.load_dir:
@@ -187,8 +200,8 @@ class Trainer:
if step != 0 and self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0:
self.save_checkpoint(self.config.graph.model_dir, step)
if step % self.config.graph.steps_per_test == 0:
self.test_image(image_idx=0, step=step)
_ = self.test_image(image_idx=10, step=step) if not self.dry_run else None
for image_idx in self.config.data.validation_image_indices:
self.test_image(image_idx=image_idx, step=step)
self.stats.print_stats(i / num_iterations)
self.stats.update_time(Stats.TOTAL_TRAIN_TIME, train_start, time(), step=-1)
@@ -198,10 +211,14 @@ class Trainer:
def train_iteration(self, batch: dict, step: int):
"""Run one iteration with a batch of inputs."""
# move batch to correct device
ray_indices = batch["indices"].to(f"cuda:{self.local_rank}")
ray_indices = batch["indices"].to(self.device)
graph_outputs = self.graph(ray_indices)
batch["pixels"] = batch["pixels"].to(f"cuda:{self.local_rank}")
losses = self.graph.get_losses(batch, graph_outputs) # or self.graph.module
batch["pixels"] = batch["pixels"].to(self.device)
losses = (
self.graph.module.get_losses(batch, graph_outputs)
if hasattr(self.graph, "module")
else self.graph.get_losses(batch, graph_outputs)
)
loss_sum, loss_dict = self.get_aggregated_loss(losses)
self.optimizers.zero_grad_all()
loss_sum.backward()
@@ -225,7 +242,7 @@ class Trainer:
rgb_coarse = []
rgb_fine = []
for i in range(0, num_rays, chunk_size):
ray_indices = all_ray_indices[i : i + chunk_size].to(f"cuda:{self.local_rank}")
ray_indices = all_ray_indices[i : i + chunk_size].to(self.device)
graph_outputs = self.graph(ray_indices)
rgb_coarse.append(graph_outputs["rgb_coarse"])
rgb_fine.append(graph_outputs["rgb_fine"])
mattport/utils/writer.py
+4 -4
View File
@@ -31,9 +31,9 @@ def get_tensorboard_name(name: str, group: str = None, prefix: str = None):
class Writer:
"""Writer class"""
def __init__(self, local_rank: int, world_size: int, save_dir: str):
def __init__(self, is_main_thread: bool, save_dir: str):
self.is_main_thread = is_main_thread
self.save_dir = save_dir
self.is_main_thread = local_rank % world_size == 0
@abstractmethod
def write_image(
@@ -83,8 +83,8 @@ class Writer:
class TensorboardWriter(Writer):
"""Tensorboard Writer Class"""
def __init__(self, local_rank: int, world_size: int, save_dir: str):
super().__init__(local_rank, world_size, save_dir)
def __init__(self, is_main_thread: bool, save_dir: str):
super().__init__(is_main_thread, save_dir)
if self.is_main_thread:
self.tb_writer = SummaryWriter(log_dir=self.save_dir)
scripts/run_train.py
+18 -12
View File
@@ -96,7 +96,7 @@ def _distributed_worker(
assert num_gpus_per_machine <= torch.cuda.device_count()
torch.cuda.set_device(local_rank)
_set_random_seed(1234 + local_rank)
_set_random_seed(1234 + local_rank) # NOTE(ethan): will this work for multiple machines? should we use global_rank?
_synchronize(world_size, local_rank)
output = main_func(local_rank, world_size, config)
@@ -105,7 +105,7 @@ def _distributed_worker(
return output
def _train(local_rank: int, world_size: int, config: DictConfig) -> Any:
def _train(local_rank: int, world_size: int, config: DictConfig, cpu: bool = False) -> Any:
"""Main training function that sets up and runs the trainer per process
Args:
@@ -116,7 +116,7 @@ def _train(local_rank: int, world_size: int, config: DictConfig) -> Any:
Returns:
Any: TODO(): determine the return type
"""
trainer = Trainer(config, local_rank, world_size)
trainer = Trainer(config, local_rank, world_size, cpu=cpu)
trainer.setup()
trainer.train()
return 0
@@ -143,7 +143,20 @@ def launch(
timeout (datetime, optional): timeout of the distributed workers Defaults to DEFAULT_TIMEOUT.
"""
world_size = num_machines * num_gpus_per_machine
if world_size > 1:
if world_size == 0:
# Using only CPU and one process.
main_func(local_rank=0, world_size=0, config=config, cpu=True)
elif world_size == 1:
# Using one gpu and one process.
try:
main_func(local_rank=0, world_size=1, config=config)
except KeyboardInterrupt:
pass
finally:
if config.logging.enable_profiler:
profiler.PROFILER.print_profile()
elif world_size > 1:
# Using multiple gpus with multiple processes.
if dist_url == "auto":
assert num_machines == 1, "dist_url=auto not supported in multi-machine jobs."
port = _find_free_port()
@@ -169,6 +182,7 @@ def launch(
try:
process_context.join()
except KeyboardInterrupt:
logger = logging.getLogger(__name__)
for i, process in enumerate(process_context.processes):
if process.is_alive():
logger.info("Terminating process %s", str(i))
@@ -178,14 +192,6 @@ def launch(
finally:
if config.logging.enable_profiler:
profiler.PROFILER.print_profile()
else:
try:
main_func(local_rank=0, world_size=1, config=config)
except KeyboardInterrupt:
pass
finally:
if config.logging.enable_profiler:
profiler.PROFILER.print_profile()
@hydra.main(config_path="../configs", config_name="default.yaml")
tests/configs/test_default.yml
+8 -5
View File
@@ -34,18 +34,21 @@ data:
data_directory: data/blender/lego_test
dataset_type: blender
downscale_factor: 4
# validation options
validation_image_indices: [0]
# dataloader options
dataloader:
num_workers: 4
shuffle: true
num_images_to_sample_from: 1
num_rays_per_batch: 1024
num_rays_per_batch: 8
num_times_to_repeat_images: 40
# machine options
machine_config:
num_gpus: 1
num_gpus: 0
num_machines: 1
machine_rank: 0
dist_url: auto
@@ -67,8 +70,8 @@ graph:
_target_: mattport.nerf.graph.vanilla_nerf.NeRFGraph
near_plane: 1.0
far_plane: 6.0
num_coarse_samples: 64
num_importance_samples: 64
num_coarse_samples: 4
num_importance_samples: 4
# optimizer options
param_groups:
fields:
@@ -85,4 +88,4 @@ graph:
lr: 1e-4
scheduler:
type: ExponentialLR
gamma: 0.1
gamma: 0.1