From 49503c2175a71c3d71b75d5702c4ff9e8b0342cb Mon Sep 17 00:00:00 2001 From: Jose <34888496+Jerry-Master@users.noreply.github.com> Date: Fri, 1 Dec 2023 01:04:01 +0100 Subject: [PATCH 1/2] Mps fix (#2436) * mps bugfix * removed docs change * Double quotes --------- Co-authored-by: Brent Yi --- nerfstudio/engine/trainer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index fe4aa85ce..4cc296063 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -464,6 +464,7 @@ class Trainer: self.optimizers.zero_grad_all() cpu_or_cuda_str: str = self.device.split(":")[0] + cpu_or_cuda_str = "cpu" if cpu_or_cuda_str == "mps" else cpu_or_cuda_str assert ( self.gradient_accumulation_steps > 0 ), f"gradient_accumulation_steps must be > 0, not {self.gradient_accumulation_steps}" From 0cb410001b1fff668e51b12e7ff579711c1187f9 Mon Sep 17 00:00:00 2001 From: Panteleris Paschalis Date: Fri, 1 Dec 2023 03:11:11 +0200 Subject: [PATCH 2/2] Fixes bug in PairPixelSampler when working with masked dataset (#2368) * Fixes bug in PairPixelSampler when working with masked dataset * Formatting * Black formatting now ok. --------- Co-authored-by: AdamRashid96 <71362382+AdamRashid96@users.noreply.github.com> Co-authored-by: Brent Yi --- nerfstudio/data/pixel_samplers.py | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 9234a5d42..07bab5d82 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -17,13 +17,7 @@ Code for sampling pixels. """ import random - -import torch -from jaxtyping import Int -from torch import Tensor - from dataclasses import dataclass, field -from nerfstudio.data.utils.pixel_sampling_utils import erode_mask from typing import ( Dict, Optional, @@ -31,9 +25,14 @@ from typing import ( Union, ) +import torch +from jaxtyping import Int +from torch import Tensor + from nerfstudio.configs.base_config import ( InstantiateConfig, ) +from nerfstudio.data.utils.pixel_sampling_utils import erode_mask @dataclass @@ -398,19 +397,18 @@ class PairPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods device: Union[torch.device, str] = "cpu", ) -> Int[Tensor, "batch_size 3"]: rays_to_sample = self.rays_to_sample + if batch_size is not None: + assert ( + int(batch_size) % 2 == 0 + ), f"PairPixelSampler can only return batch sizes in multiples of two (got {batch_size})" + rays_to_sample = batch_size // 2 + if isinstance(mask, Tensor): m = erode_mask(mask.permute(0, 3, 1, 2).float(), pixel_radius=self.radius) nonzero_indices = torch.nonzero(m[:, 0], as_tuple=False).to(device) chosen_indices = random.sample(range(len(nonzero_indices)), k=rays_to_sample) indices = nonzero_indices[chosen_indices] else: - rays_to_sample = self.rays_to_sample - if batch_size is not None: - assert ( - int(batch_size) % 2 == 0 - ), f"PairPixelSampler can only return batch sizes in multiples of two (got {batch_size})" - rays_to_sample = batch_size // 2 - s = (rays_to_sample, 1) ns = torch.randint(0, num_images, s, dtype=torch.long, device=device) hs = torch.randint(self.radius, image_height - self.radius, s, dtype=torch.long, device=device)