From 22fdcfe235e5dcc98c1a8f12ac396ab42d26c7e4 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Wed, 18 May 2022 17:45:31 -0700 Subject: [PATCH 1/4] Add occupancy to samplers --- mattport/nerf/fields/base.py | 10 +-- mattport/nerf/graph/vanilla_nerf.py | 3 +- mattport/nerf/sampler.py | 121 +++++++++++----------------- mattport/structures/rays.py | 12 ++- 4 files changed, 62 insertions(+), 84 deletions(-) diff --git a/mattport/nerf/fields/base.py b/mattport/nerf/fields/base.py index 5f3b01f37..e930b2174 100644 --- a/mattport/nerf/fields/base.py +++ b/mattport/nerf/fields/base.py @@ -25,23 +25,23 @@ class Field(nn.Module): return density @abstractmethod - def get_density(self, point_samples: PointSamples, valid_mask=None): + def get_density(self, point_samples: PointSamples): """Computes and returns the densities.""" @abstractmethod def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): """Computes and returns the colors.""" - def forward(self, point_samples: PointSamples, valid_mask_density=None): + def forward(self, point_samples: PointSamples): """Evaluates the field at points along the ray.""" - density, density_embedding = self.get_density(point_samples, valid_mask=valid_mask_density) + density, density_embedding = self.get_density(point_samples) field_outputs = self.get_outputs(point_samples, density_embedding=density_embedding) field_outputs["density"] = density return field_outputs - def forward_with_weight_pruning(self, ray_samples: RaySamples, valid_mask_density=None): + def forward_with_weight_pruning(self, ray_samples: RaySamples): """Forward but with pruning based on weights, from the densities.""" - density, density_embedding = self.get_density(ray_samples, valid_mask=valid_mask_density) + density, density_embedding = self.get_density(ray_samples) weights = ray_samples.get_weights(density) valid_mask_outputs = weights >= self.weights_threshold field_outputs = self.get_outputs( diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index ffb90bea2..931e2bec2 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -81,9 +81,10 @@ class NeRFField(Field): self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_rgb.get_out_dim()) self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) - def get_density(self, point_samples: PointSamples, valid_mask=None): + def get_density(self, point_samples: PointSamples): """Computes and returns the densities.""" positions = point_samples.positions + valid_mask = point_samples.valid_mask if not is_not_none(valid_mask): valid_mask = torch.ones_like(positions[..., 0]).bool() # placeholders for values to return diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index 0ddea9091..e3e4cb091 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -3,98 +3,43 @@ Collection of sampling strategies """ from typing import Optional + import torch from torch import nn from torchtyping import TensorType + from mattport.nerf.occupancy_grid import OccupancyGrid - from mattport.structures.rays import RayBundle, RaySamples -class OccupancyGridSampler(nn.Module): - """Sample based on occupancy grid""" - - def __init__( - self, - occupancy_grid: OccupancyGrid, - num_samples: int = 256, - weight_threshold: float = 1e-4, - train_stratified=True, - ) -> None: - """ - Args: - """ - super().__init__() - self.num_samples = num_samples - self.train_stratified = train_stratified - self.weight_threshold = weight_threshold - self.occupancy_grid = occupancy_grid - - @torch.no_grad() - def forward( - self, - ray_bundle: RayBundle, - num_samples: Optional[int] = None, - ) -> RaySamples: - """Generates position samples uniformly. - - Args: - ray_bundle (RayBundle): Rays to generate samples for - num_samples (Optional[int]): Number of samples per ray - - Returns: - RaySamples: Positions and deltas for samples along a ray - """ - assert ray_bundle.nears is not None - assert ray_bundle.fars is not None - - num_samples = num_samples or self.num_samples - num_rays = ray_bundle.origins.shape[0] - - bins = torch.linspace(0.0, 1.0, num_samples + 1).to(ray_bundle.origins.device) # shape (num_samples+1,) - bins = ray_bundle.nears[:, None] + bins[None, :] * ( - ray_bundle.fars[:, None] - ray_bundle.nears[:, None] - ) # shape (num_rays, num_samples+1) - - if self.train_stratified and self.training: - t_rand = torch.rand((num_rays, num_samples), dtype=bins.dtype, device=bins.device) - ts = bins[:, :-1] + t_rand * (bins[:, 1:] - bins[:, :-1]) # shape (num_rays, num_samples) - else: - ts = (bins[:, 1:] + bins[:, :-1]) / 2 - - ray_samples = RaySamples( - ts=ts, - ray_bundle=ray_bundle, - ) - - densities = self.occupancy_grid.get_densities(ray_samples.positions, update_iter_count=True) - weights = ray_samples.get_weights(densities) - - valid_mask = weights >= self.weight_threshold - - return ray_samples, weights, valid_mask - - class UniformSampler(nn.Module): """Sample uniformly along a ray""" - def __init__(self, num_samples: int, train_stratified=True) -> None: + def __init__( + self, + num_samples: int, + train_stratified=True, + occupancy_field: Optional[OccupancyGrid] = None, + weight_threshold: float = 1e-4, + ) -> None: """ Args: - near_plane (float): Minimum distance along ray to sample - far_plane (float): Maximum distance along ray to sample num_samples (int): Number of samples per ray train_stratified (bool): Use stratified sampling during training. Defults to True + occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides, + samples below weight_threshold as set as invalid. + weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. """ super().__init__() - self.num_samples = num_samples self.train_stratified = train_stratified + self.occupancy_field = occupancy_field + self.weight_threshold = weight_threshold @torch.no_grad() def forward( self, - ray_bundle: RayBundle, + ray_bundle: RayBundle = None, num_samples: Optional[int] = None, ) -> RaySamples: """Generates position samples uniformly. @@ -128,29 +73,48 @@ class UniformSampler(nn.Module): ray_bundle=ray_bundle, ) + if self.occupancy_field is not None: + densities = self.occupancy_field.get_densities(ray_samples.positions, update_iter_count=True) + weights = ray_samples.get_weights(densities) + + valid_mask = weights >= self.weight_threshold + ray_samples.set_valid_mask(valid_mask) + return ray_samples class PDFSampler(nn.Module): """Sample based on probability distribution""" - def __init__(self, num_samples: int, train_stratified: bool = True, include_original: bool = True) -> None: + def __init__( + self, + num_samples: int, + train_stratified: bool = True, + include_original: bool = True, + occupancy_field: OccupancyGrid = None, + weight_threshold: float = 1e-4, + ) -> None: """ Args: num_samples (int): Number of samples per ray train_stratified: boolean: Randomize location within each bin during training. Defaults to True include_original: Add original samples to ray. Defaults to True + occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides, + samples below weight_threshold as set as invalid. + weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. """ super().__init__() self.num_samples = num_samples self.include_original = include_original self.train_stratified = train_stratified + self.occupancy_field = occupancy_field + self.weight_threshold = weight_threshold @torch.no_grad() def forward( self, - coarse_ray_samples: RaySamples, + ray_samples: RaySamples, weights: TensorType[..., "num_samples"], num_samples: Optional[int] = None, eps: float = 1e-5, @@ -195,7 +159,7 @@ class PDFSampler(nn.Module): indicies_g = torch.stack([below, above], -1) matched_shape = (indicies_g.shape[0], indicies_g.shape[1], cdf.shape[-1]) cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), dim=2, index=indicies_g) - bins_g = torch.gather(coarse_ray_samples.ts.unsqueeze(1).expand(matched_shape), dim=2, index=indicies_g) + bins_g = torch.gather(ray_samples.ts.unsqueeze(1).expand(matched_shape), dim=2, index=indicies_g) denom = cdf_g[..., 1] - cdf_g[..., 0] denom = torch.where(denom < eps, torch.ones_like(denom), denom) @@ -203,7 +167,7 @@ class PDFSampler(nn.Module): ts = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0]) if self.include_original: - ts, _ = torch.sort(torch.cat([coarse_ray_samples.ts, ts], -1), -1) + ts, _ = torch.sort(torch.cat([ray_samples.ts, ts], -1), -1) else: ts, _ = torch.sort(ts, -1) @@ -212,7 +176,14 @@ class PDFSampler(nn.Module): ray_samples = RaySamples( ts=ts, - ray_bundle=coarse_ray_samples.ray_bundle, + ray_bundle=ray_samples.ray_bundle, ) + if self.occupancy_field is not None: + densities = self.occupancy_field.get_densities(ray_samples.positions, update_iter_count=True) + weights = ray_samples.get_weights(densities) + + valid_mask = weights >= self.weight_threshold + ray_samples.set_valid_mask(valid_mask) + return ray_samples diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index d1df55b97..56a8134f7 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -4,8 +4,8 @@ Some ray datastructures. import random from dataclasses import dataclass from typing import Optional -import torch +import torch from torchtyping import TensorType from mattport.utils.misc import is_not_none @@ -18,6 +18,7 @@ class PointSamples: positions: TensorType[..., 3] = None directions: TensorType[..., 3] = None camera_indices: TensorType[..., 1] = None + valid_mask: TensorType[..., 1] = None @dataclass @@ -145,11 +146,12 @@ class RaySamples: self.positions = self.get_positions(ray_bundle) self.directions = ray_bundle.directions.unsqueeze(1).repeat(1, self.positions.shape[1], 1) self.deltas = self.get_deltas() + self.valid_mask = torch.ones_like(self.ts, dtype=torch.bool) - def to_point_samples(self): + def to_point_samples(self) -> PointSamples: """Convert to PointSamples instance and return.""" # TODO: make this more interpretable - return PointSamples(positions=self.positions, directions=self.directions) + return PointSamples(positions=self.positions, directions=self.directions, valid_mask=self.valid_mask) def get_camera_indices(self): """Returns camera indices.""" @@ -191,3 +193,7 @@ class RaySamples: weights = alphas * transmittance # [..., "num_samples"] return weights + + def set_valid_mask(self, valid_mask: TensorType[..., "num_samples"]) -> None: + """Sets valid mask""" + self.valid_mask = valid_mask From 998b98b19a1ecbae7653323bc81b75021e22df80 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Thu, 19 May 2022 14:17:03 -0700 Subject: [PATCH 2/4] Clean up sampler --- mattport/nerf/sampler.py | 68 ++++++++++++++++++++++------------------ 1 file changed, 38 insertions(+), 30 deletions(-) diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index e3e4cb091..9269322d9 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -2,6 +2,7 @@ Collection of sampling strategies """ +from abc import abstractmethod from typing import Optional import torch @@ -12,7 +13,34 @@ from mattport.nerf.occupancy_grid import OccupancyGrid from mattport.structures.rays import RayBundle, RaySamples -class UniformSampler(nn.Module): +class Sampler(nn.Module): + """Generate Samples""" + + def __init__( + self, num_samples: int, occupancy_field: Optional[OccupancyGrid] = None, weight_threshold: float = 1e-4 + ) -> None: + super().__init__() + self.num_samples = num_samples + self.occupancy_field = occupancy_field + self.weight_threshold = weight_threshold + + @abstractmethod + def generate_ray_samples(self) -> RaySamples: + """Generate Ray Samples""" + + def forward(self, *args, **kwargs) -> RaySamples: + """Generate ray samples with optional occupancy filtering""" + ray_samples = self.generate_ray_samples(*args, **kwargs) + if self.occupancy_field is not None: + densities = self.occupancy_field.get_densities(ray_samples.positions, update_iter_count=True) + weights = ray_samples.get_weights(densities) + + valid_mask = weights >= self.weight_threshold + ray_samples.set_valid_mask(valid_mask) + return ray_samples + + +class UniformSampler(Sampler): """Sample uniformly along a ray""" def __init__( @@ -30,14 +58,11 @@ class UniformSampler(nn.Module): samples below weight_threshold as set as invalid. weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. """ - super().__init__() - self.num_samples = num_samples + super().__init__(num_samples=num_samples, occupancy_field=occupancy_field, weight_threshold=weight_threshold) self.train_stratified = train_stratified - self.occupancy_field = occupancy_field - self.weight_threshold = weight_threshold @torch.no_grad() - def forward( + def generate_ray_samples( self, ray_bundle: RayBundle = None, num_samples: Optional[int] = None, @@ -51,6 +76,7 @@ class UniformSampler(nn.Module): Returns: RaySamples: Positions and deltas for samples along a ray """ + assert ray_bundle is not None assert ray_bundle.nears is not None assert ray_bundle.fars is not None @@ -73,17 +99,10 @@ class UniformSampler(nn.Module): ray_bundle=ray_bundle, ) - if self.occupancy_field is not None: - densities = self.occupancy_field.get_densities(ray_samples.positions, update_iter_count=True) - weights = ray_samples.get_weights(densities) - - valid_mask = weights >= self.weight_threshold - ray_samples.set_valid_mask(valid_mask) - return ray_samples -class PDFSampler(nn.Module): +class PDFSampler(Sampler): """Sample based on probability distribution""" def __init__( @@ -103,19 +122,15 @@ class PDFSampler(nn.Module): samples below weight_threshold as set as invalid. weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. """ - super().__init__() - - self.num_samples = num_samples - self.include_original = include_original + super().__init__(num_samples=num_samples, occupancy_field=occupancy_field, weight_threshold=weight_threshold) self.train_stratified = train_stratified - self.occupancy_field = occupancy_field - self.weight_threshold = weight_threshold + self.include_original = include_original @torch.no_grad() - def forward( + def generate_ray_samples( self, - ray_samples: RaySamples, - weights: TensorType[..., "num_samples"], + ray_samples: RaySamples = None, + weights: TensorType[..., "num_samples"] = None, num_samples: Optional[int] = None, eps: float = 1e-5, ) -> RaySamples: @@ -179,11 +194,4 @@ class PDFSampler(nn.Module): ray_bundle=ray_samples.ray_bundle, ) - if self.occupancy_field is not None: - densities = self.occupancy_field.get_densities(ray_samples.positions, update_iter_count=True) - weights = ray_samples.get_weights(densities) - - valid_mask = weights >= self.weight_threshold - ray_samples.set_valid_mask(valid_mask) - return ray_samples From 8a6b7c3e465d272d7aa940c6c60e73f82f3d18c2 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Thu, 19 May 2022 18:38:30 -0700 Subject: [PATCH 3/4] Refactor ray samples --- mattport/nerf/graph/vanilla_nerf.py | 2 +- mattport/nerf/sampler.py | 11 +-- mattport/structures/rays.py | 142 +++++++++++++++------------- tests/nerf/test_sampler.py | 2 +- 4 files changed, 81 insertions(+), 76 deletions(-) diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 931e2bec2..5eb4f7fe7 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -179,7 +179,7 @@ class NeRFGraph(Graph): depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform.ts) # pdf sampling - ray_samples_pdf = self.sampler_pdf(ray_samples_uniform, weights_coarse) + ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) # fine field: field_outputs_fine = self.field_fine.forward(ray_samples_pdf.to_point_samples()) diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index 9269322d9..33a85921d 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -94,10 +94,7 @@ class UniformSampler(Sampler): else: ts = (bins[:, 1:] + bins[:, :-1]) / 2 - ray_samples = RaySamples( - ts=ts, - ray_bundle=ray_bundle, - ) + ray_samples = ray_bundle.get_ray_samples(ts) return ray_samples @@ -129,6 +126,7 @@ class PDFSampler(Sampler): @torch.no_grad() def generate_ray_samples( self, + ray_bundle: RayBundle = None, ray_samples: RaySamples = None, weights: TensorType[..., "num_samples"] = None, num_samples: Optional[int] = None, @@ -189,9 +187,6 @@ class PDFSampler(Sampler): # Stop gradients ts = ts.detach() - ray_samples = RaySamples( - ts=ts, - ray_bundle=ray_samples.ray_bundle, - ) + ray_samples = ray_bundle.get_ray_samples(ts) return ray_samples diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index 56a8134f7..7593041ea 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -21,6 +21,50 @@ class PointSamples: valid_mask: TensorType[..., 1] = None +@dataclass +class RaySamples: + """Samples along a ray""" + + positions: TensorType[..., 3] = None + directions: TensorType[..., 3] = None + camera_indices: TensorType[..., 1] = None + valid_mask: TensorType[..., 1] = None + ts: TensorType[..., 1] = None + deltas: TensorType[..., 1] = None + + def to_point_samples(self) -> PointSamples: + """Convert to PointSamples instance and return.""" + # TODO: make this more interpretable + return PointSamples(positions=self.positions, directions=self.directions, valid_mask=self.valid_mask) + + def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]: + """Return weights based on predicted densities + + Args: + densities (TensorType[..., "num_samples", 1]): Predicted densities for samples along ray + + Returns: + TensorType[..., "num_samples"]: Weights for each sample + """ + + delta_density = self.deltas * densities[..., 0] + alphas = 1 - torch.exp(-delta_density) + + transmittance = torch.cumsum(delta_density[..., :-1], dim=-1) + transmittance = torch.cat( + [torch.zeros((*transmittance.shape[:1], 1)).to(densities.device), transmittance], axis=-1 + ) + transmittance = torch.exp(-transmittance) # [..., "num_samples"] + + weights = alphas * transmittance # [..., "num_samples"] + + return weights + + def set_valid_mask(self, valid_mask: TensorType[..., "num_samples"]) -> None: + """Sets valid mask""" + self.valid_mask = valid_mask + + @dataclass class RayBundle: """_summary_ @@ -86,6 +130,38 @@ class RayBundle: valid_mask=self.valid_mask[valid_mask] if is_not_none(self.valid_mask) else None, ) + def get_ray_samples(self, ts: TensorType["num_rays", "num_samples"]) -> RaySamples: + """ + Args: + ts (TensorType["num_rays", "num_samples"]): _description_ + + Returns: + RaySamples: _description_ + """ + positions = self.origins[:, None] + ts[:, :, None] * self.directions[:, None] + directions = self.directions.unsqueeze(1).repeat(1, positions.shape[1], 1) + valid_mask = torch.ones_like(ts, dtype=torch.bool) + + dists = ts[..., 1:] - ts[..., :-1] + dists = torch.cat([dists, dists[..., -1:]], -1) # [N_rays, N_samples] + deltas = dists * torch.norm(self.directions[..., None, :], dim=-1) + + if self.camera_indices is not None: + camera_indices = self.camera_indices.unsqueeze(1).repeat(1, positions.shape[1]) + else: + camera_indices = None + + ray_samples = RaySamples( + positions=positions, + directions=directions, + camera_indices=camera_indices, + valid_mask=valid_mask, + ts=ts, + deltas=deltas, + ) + + return ray_samples + @dataclass class CameraRayBundle: @@ -131,69 +207,3 @@ class CameraRayBundle: directions=self.directions.view(-1, 3)[start_idx:end_idx], camera_indices=camera_indices, ) - - -class RaySamples: - """_summary_""" - - def __init__( - self, - ts: TensorType["num_rays", "num_samples"], - ray_bundle: RayBundle, - ) -> None: - self.ray_bundle = ray_bundle - self.ts = ts - self.positions = self.get_positions(ray_bundle) - self.directions = ray_bundle.directions.unsqueeze(1).repeat(1, self.positions.shape[1], 1) - self.deltas = self.get_deltas() - self.valid_mask = torch.ones_like(self.ts, dtype=torch.bool) - - def to_point_samples(self) -> PointSamples: - """Convert to PointSamples instance and return.""" - # TODO: make this more interpretable - return PointSamples(positions=self.positions, directions=self.directions, valid_mask=self.valid_mask) - - def get_camera_indices(self): - """Returns camera indices.""" - assert not isinstance(self.ray_bundle.camera_indices, type(None)), "Camera indices cannot be None here." - camera_indices = self.ray_bundle.camera_indices.unsqueeze(1).repeat(1, self.positions.shape[1]) - return camera_indices - - def get_positions(self, ray_bundle: RayBundle) -> TensorType["num_rays", "num_samples", 3]: - """Returns positions.""" - return ray_bundle.origins[:, None] + self.ts[:, :, None] * ray_bundle.directions[:, None] - - def get_deltas(self) -> TensorType[..., "num_samples"]: - """Returns deltas.""" - dists = self.ts[..., 1:] - self.ts[..., :-1] - dist_inf = 1e10 * torch.ones_like(dists[..., -1:]) - dists = torch.cat([dists, dist_inf], -1) # [N_rays, N_samples] - deltas = dists * torch.norm(self.ray_bundle.directions[..., None, :], dim=-1) - return deltas - - def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]: - """Return weights based on predicted densities - - Args: - densities (TensorType[..., "num_samples", 1]): Predicted densities for samples along ray - - Returns: - TensorType[..., "num_samples"]: Weights for each sample - """ - - delta_density = self.deltas * densities[..., 0] - alphas = 1 - torch.exp(-delta_density) - - transmittance = torch.cumsum(delta_density[..., :-1], dim=-1) - transmittance = torch.cat( - [torch.zeros((*transmittance.shape[:1], 1)).to(densities.device), transmittance], axis=-1 - ) - transmittance = torch.exp(-transmittance) # [..., "num_samples"] - - weights = alphas * transmittance # [..., "num_samples"] - - return weights - - def set_valid_mask(self, valid_mask: TensorType[..., "num_samples"]) -> None: - """Sets valid mask""" - self.valid_mask = valid_mask diff --git a/tests/nerf/test_sampler.py b/tests/nerf/test_sampler.py index 92a6848bb..df88b889c 100644 --- a/tests/nerf/test_sampler.py +++ b/tests/nerf/test_sampler.py @@ -47,7 +47,7 @@ def test_pdf_sampler(): # Just check that it doesn't crash. Need to add some actual tests. pdf_sampler = PDFSampler(num_samples) - pdf_sampler(coarse_ray_samples, weights, num_samples) + pdf_sampler(ray_bundle, coarse_ray_samples, weights, num_samples) # TODO Tancik: Add more precise tests From 10d713f84b00f4ea07adc20ef66d3342cfefc282 Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 20 May 2022 08:51:45 -0700 Subject: [PATCH 4/4] Consolidate StatsTracker as LocalWriter; cleaned up functions to write (#88) * removed stats tracker; consolidate writers into local writer * basic revamp of stats tracker/writers --- configs/logging/default_logging.yaml | 10 +- mattport/nerf/graph/nerfw.py | 2 +- mattport/nerf/graph/semantic_nerf.py | 2 +- mattport/nerf/graph/vanilla_nerf.py | 20 +-- mattport/nerf/optimizers.py | 2 +- mattport/nerf/trainer.py | 60 ++++--- mattport/utils/decorators.py | 34 +--- mattport/utils/stats_tracker.py | 192 --------------------- mattport/utils/writer.py | 245 ++++++++++++++++++++++----- 9 files changed, 256 insertions(+), 311 deletions(-) delete mode 100644 mattport/utils/stats_tracker.py diff --git a/configs/logging/default_logging.yaml b/configs/logging/default_logging.yaml index 59ef1333b..8938d1233 100644 --- a/configs/logging/default_logging.yaml +++ b/configs/logging/default_logging.yaml @@ -1,12 +1,12 @@ steps_per_log: 10 -# stats tracker logs configurable training stats per iteration to terminal (e.g. rays/s, PSNR) -enable_stats: True -stats_tracker: - max_history: 10 - stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] # see mattport/utils/stats_tracker.py for options +max_history: 20 # writer logs losses and images per iteration writer: TensorboardWriter: save_dir: './' + LocalWriter: + save_dir: './' + stats_to_track: ["Data Load (time)", "Train Iter (time)", "Rays Per Sec (time)", "Test PSNR"] + # profiler logs run times of functions and prints at end of training enable_profiler: True \ No newline at end of file diff --git a/mattport/nerf/graph/nerfw.py b/mattport/nerf/graph/nerfw.py index 06ae5d81f..b31061266 100644 --- a/mattport/nerf/graph/nerfw.py +++ b/mattport/nerf/graph/nerfw.py @@ -212,4 +212,4 @@ class NerfWGraph(NeRFGraph): row2 = torch.cat([depth_fine, depth_fine_static, depth_coarse], dim=-2) combined_image = torch.cat([row0, row1, row2], dim=-3) - writer.write_image(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img") + writer.put_scalar(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img") diff --git a/mattport/nerf/graph/semantic_nerf.py b/mattport/nerf/graph/semantic_nerf.py index 02357caf6..726c54d77 100644 --- a/mattport/nerf/graph/semantic_nerf.py +++ b/mattport/nerf/graph/semantic_nerf.py @@ -155,4 +155,4 @@ class SemanticNerfGraph(NeRFGraph): self.stuff_classes ) semantic_labels_image = visualization.apply_colormap(semantic_labels[..., None]) - writer.write_image(name="image_idx_{image_idx}", image=semantic_labels_image, step=step, group="semantics") + writer.put_image(name="image_idx_{image_idx}", image=semantic_labels_image, step=step, group="semantics") diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 5eb4f7fe7..f1ba247ba 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -23,7 +23,7 @@ from mattport.nerf.sampler import PDFSampler, UniformSampler from mattport.structures import colors from mattport.structures.rays import PointSamples, RayBundle from mattport.nerf.colliders import NearFarCollider -from mattport.utils import stats_tracker, visualization, writer +from mattport.utils import visualization, writer from mattport.utils.misc import is_not_none @@ -232,9 +232,9 @@ class NeRFGraph(Graph): combined_acc = torch.cat([acc_coarse, acc_fine], dim=1) combined_depth = torch.cat([depth_coarse, depth_fine], dim=1) - writer.write_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img") - writer.write_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation") - writer.write_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth") + writer.put_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img") + writer.put_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation") + writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth") # Switch images from [H, W, C] to [1, C, H, W] for metrics computations image = torch.moveaxis(image, -1, 0)[None, ...] @@ -246,13 +246,11 @@ class NeRFGraph(Graph): fine_ssim = self.ssim(image, rgb_fine) fine_lpips = self.lpips(image, rgb_fine) - writer.write_scalar(name=f"val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step, group="psnr") - writer.write_scalar(name=f"val_{image_idx}-fine", scalar=float(fine_psnr), step=step, group="psnr") - writer.write_scalar(name=f"val_{image_idx}", scalar=float(fine_ssim), step=step, group="ssim") - writer.write_scalar(name=f"val_{image_idx}", scalar=float(fine_lpips), step=step, group="lpips") + writer.put_scalar(name=f"val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step, group="psnr") + writer.put_scalar(name=f"val_{image_idx}-fine", scalar=float(fine_psnr), step=step, group="psnr") + writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_ssim), step=step, group="ssim") + writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_lpips), step=step, group="lpips") - stats_tracker.update_stats( - {"name": stats_tracker.Stats.CURR_TEST_PSNR, "value": float(fine_psnr), "step": step} - ) + writer.put_scalar(name="Test PSNR", scalar=float(fine_psnr), step=step) return fine_psnr.item() diff --git a/mattport/nerf/optimizers.py b/mattport/nerf/optimizers.py index 917eeed0d..181ac1517 100644 --- a/mattport/nerf/optimizers.py +++ b/mattport/nerf/optimizers.py @@ -83,4 +83,4 @@ class Optimizers: scheduler.step() # TODO(ethan): clean this up. why is there indexing into a list? lr = scheduler.get_last_lr()[0] - writer.write_scalar(name=f"{param_group_name}", scalar=lr, step=step, group="learning_rate") + writer.put_scalar(name=f"{param_group_name}", scalar=lr, step=step, group="learning_rate") diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index ec2486a92..e6f654a24 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -18,7 +18,7 @@ from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch from mattport.nerf.dataset.utils import DatasetInputs, get_dataset_inputs from mattport.nerf.optimizers import Optimizers -from mattport.utils import profiler, stats_tracker, writer +from mattport.utils import profiler, writer from mattport.utils.decorators import check_main_thread logging.getLogger("PIL").setLevel(logging.WARNING) @@ -51,8 +51,7 @@ class Trainer: self.optimizers = None self.start_step = 0 # logging variables - stats_tracker.setup_stats_tracker(config) - writer.setup_event_writers(config.logging.writer) + writer.setup_event_writers(config) profiler.setup_profiler(config.logging) self.device = "cpu" if self.world_size == 0 else f"cuda:{self.local_rank}" @@ -199,47 +198,54 @@ class Trainer: train_start = time() num_iterations = self.config.graph.max_num_iterations iter_dataset = iter(self.train_dataloader) - for i, step in enumerate(range(self.start_step, self.start_step + num_iterations)): + for step in range(self.start_step, self.start_step + num_iterations): data_start = time() batch = next(iter_dataset) - stats_tracker.update_stats( - {"name": stats_tracker.Stats.ITER_LOAD_TIME, "start_time": data_start, "end_time": time(), "step": step} + writer.put_time( + name="Data Load (time)", start_time=data_start, end_time=time(), step=step, avg_over_iters=True ) iter_start = time() loss_dict = self.train_iteration(batch, step) - stats_tracker.update_stats( - { - "name": stats_tracker.Stats.RAYS_PER_SEC, - "start_time": iter_start, - "end_time": time(), - "step": step, - "batch_size": batch["indices"].shape[0], - }, + writer.put_time( + name="Rays Per Sec (time)", + start_time=iter_start, + end_time=time(), + step=step, + avg_over_iters=True, + avg_over_batch=batch["indices"].shape[0], ) - stats_tracker.update_stats( - { - "name": stats_tracker.Stats.ITER_TRAIN_TIME, - "start_time": iter_start, - "end_time": time(), - "step": step, - } + writer.put_time( + name="Train Iter (time)", + start_time=iter_start, + end_time=time(), + step=step, + avg_over_iters=True, + update_eta=True, ) if step != 0 and step % self.config.logging.steps_per_log == 0: - writer.write_scalar_dict(scalar_dict=loss_dict, step=step, group="Loss", prefix="train-") + writer.put_dict(name="loss_dict", scalar_dict=loss_dict, step=step, group="Loss", prefix="train-") # TODO: add the learning rates to tensorboard/logging 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: # NOTE(ethan): we should still run this in dry-run mode! for image_idx in self.config.data.val_image_indices: _ = self.test_image(image_idx=image_idx, step=step) - stats_tracker.print_stats(i / num_iterations) + self._write_out_storage(step) - stats_tracker.update_stats( - {"name": stats_tracker.Stats.TOTAL_TRAIN_TIME, "start_time": train_start, "end_time": time()} - ) - stats_tracker.print_stats(1.0) + writer.put_time(name="Train Total (time)", start_time=train_start, end_time=time(), step=num_iterations) + self._write_out_storage(num_iterations) + + def _write_out_storage(self, step): + """Perform writes only during appropriate time steps""" + if ( + step % self.config.logging.steps_per_log == 0 + or (self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0) + or step % self.config.graph.steps_per_test == 0 + or step == self.config.graph.max_num_iterations + ): + writer.write_out_storage() @profiler.time_function def train_iteration(self, batch: dict, step: int): diff --git a/mattport/utils/decorators.py b/mattport/utils/decorators.py index f7b79f7bf..9d99c68f9 100644 --- a/mattport/utils/decorators.py +++ b/mattport/utils/decorators.py @@ -19,24 +19,12 @@ def decorate_all(decorators: List[Callable]) -> Callable: return decorate -def check_stats_enabled(func: Callable) -> Callable: - """Decorator: check if stats tracker is enabled""" - - def wrapper(self, *args, **kwargs): - ret = None - if self.config.logging.enable_stats: - ret = func(self, *args, **kwargs) - return ret - - return wrapper - - def check_profiler_enabled(func: Callable) -> Callable: """Decorator: check if profiler is enabled""" def wrapper(self, *args, **kwargs): ret = None - if self.config.enable_stats: + if self.config.enable_profiler: ret = func(self, *args, **kwargs) return ret @@ -46,26 +34,10 @@ def check_profiler_enabled(func: Callable) -> Callable: def check_main_thread(func: Callable) -> Callable: """Decorator: check if you are on main thread""" - def wrapper(self, *args, **kwargs): + def wrapper(*args, **kwargs): ret = None if comms.is_main_process(): - ret = func(self, *args, **kwargs) - return ret - - return wrapper - - -def check_print_stats_step(func: Callable) -> Callable: - """Decorator: check if it is time to print stats update""" - - def wrapper(self, *args, **kwargs): - ret = None - if ( - self.step % self.config.logging.steps_per_log == 0 - or (self.config.graph.steps_per_save and self.step % self.config.graph.steps_per_save == 0) - or self.step % self.config.graph.steps_per_test == 0 - ): - ret = func(self, *args, **kwargs) + ret = func(*args, **kwargs) return ret return wrapper diff --git a/mattport/utils/stats_tracker.py b/mattport/utils/stats_tracker.py deleted file mode 100644 index 64f0fc5f5..000000000 --- a/mattport/utils/stats_tracker.py +++ /dev/null @@ -1,192 +0,0 @@ -""" -Stats tracker base class and functionality -""" -import datetime -import enum -import logging -from typing import Any, Dict - -from omegaconf import DictConfig - -from mattport.utils import comms -from mattport.utils.decorators import check_main_thread, check_print_stats_step, check_stats_enabled, decorate_all - -STATS_TRACKER = [] - - -def update_stats(args: Dict[str, Any]): - """update global stats tracker based on arguments""" - assert comms.is_main_process(), "Writing out with process other than main" - if "value" in args: - STATS_TRACKER[0].update_value(**args) - elif "start_time" in args: - STATS_TRACKER[0].update_time(**args) - else: - raise NotImplementedError - - -def print_stats(fraction_done: float): - """print current statistics information""" - assert comms.is_main_process(), "Writing out with process other than main" - STATS_TRACKER[0].print_stats(fraction_done) - - -def setup_stats_tracker(config: DictConfig): - """Initialization of stats tracker""" - if comms.is_main_process(): - STATS_TRACKER.append(StatsTracker(config)) - - -class Stats(enum.Enum): - """Possible Stats values for StatsTracker - - Args: - enum (_type_): _description_ - """ - - ITER_LOAD_TIME = "Data Load (ms)" - ITER_TRAIN_TIME = "Train Iter (ms)" - TOTAL_TRAIN_TIME = "Train Total (time)" - RAYS_PER_SEC = "Rays Per Sec (1/s)" - CURR_TEST_PSNR = "Test PSNR" - ETA = "ETA (time)" - - -@decorate_all([check_stats_enabled, check_main_thread]) -class StatsTracker: - """Stats Tracker class""" - - def __init__(self, config: DictConfig): - self.config = config - if self.config.logging.enable_stats: - self.max_history = self.config.logging.stats_tracker.max_history - self.step = 0 - self.stats_dict = {} - self.past_stats = [] - self.new_key = False - self.stats_to_track = set() - for name in self.config.logging.stats_tracker.stats_to_track: - self.stats_to_track.add(Stats[name]) - logging.info("Successfully set up StatsTracker") - else: - logging.info("StatsTracker disabled; enable in config.") - - def update_value(self, name: enum.Enum, value: float, step: int): - """update stats dictionary with key value pair - - Args: - name (enum.Enum): Enum name of statistic we are logging - value (float): value to update. - step (int): number of total iteration steps. - """ - if name in self.stats_to_track: - self.step = step - self.new_key = not name in self.stats_dict or self.new_key - self.stats_dict[name] = value - - def update_time( - self, name: enum.Enum, start_time: float, end_time: float, step: int = None, batch_size: int = None - ): - """update the stats dictionary with running averages/cumulative durations - - Args: - name (enum.Enum): Enum name of statistic we are logging - start_time (float): start time for the call in seconds - end_time (float): end time when the call finished executing in seconds - step (int): number of total iteration steps. Defaults to None. - if None, reports duration without averaging - batch_size (int, optional): total number of rays in a batch; - if None, reports duration instead of batch per second. Defaults to None. - """ - if name in self.stats_to_track: - self.step = step - self.new_key = not name in self.stats_dict or self.new_key - val = end_time - start_time - if batch_size: - # calculate the batch per second stat - val = batch_size / val - - if step is not None: - # calculate updated average over the buffered max history - curr_history = self.stats_dict.get(name, {"buffer": [], "avg": 0}) - curr_buffer = curr_history["buffer"] - curr_avg = curr_history["avg"] - if len(curr_buffer) >= self.max_history: - curr_buffer.pop(0) - curr_buffer.append(val) - curr_avg = sum(curr_buffer) / len(curr_buffer) - self.stats_dict[name] = {"buffer": curr_buffer, "avg": curr_avg} - else: - # logging total time instead of average - self.stats_dict[name] = val - - if name == Stats.ITER_TRAIN_TIME and Stats.ETA in self.stats_to_track: - # update ETA if logging iteration train time - remain_iter = self.config.graph.max_num_iterations - step - self.stats_dict[Stats.ETA] = remain_iter * self.stats_dict[name]["avg"] - - def handle_header(self): - """helper to handle the printing of the header labels""" - if self.step == 0 or self.new_key: - mssg = f"{'Step (% Done)':<20}" - for k in self.stats_dict: - mssg += f"{k.value:<20} " - if self.step > 0: - cursorup(len(self.past_stats) + 2) - print(mssg) - print("-" * len(mssg)) - if self.step > 0: - for mssg in self.past_stats: - print(mssg) - - def handle_stats(self, fraction_done): - """helper to print out the stats in a readable format""" - if self.step > 0: - curr_mssg = f"{self.step} ({fraction_done*100:.02f}%)" - curr_mssg = f"{curr_mssg:<20}" - for k, v in self.stats_dict.items(): - if isinstance(v, dict): - v = v["avg"] - - if "(time)" in k.value: - v = str(datetime.timedelta(seconds=v)) - elif "(ms)" in k.value: - ms = v * 1e3 - v = f"{ms:0.4f}" - else: - v = f"{v:0.4f}" - curr_mssg += f"{v:<20} " - # update the history buffer - # if len(self.past_stats) >= self.max_history: - # self.past_stats.pop(0) - # cursorup(len(self.past_stats) + 1) - # for mssg in self.past_stats: - # print(mssg) - print(curr_mssg) - # self.past_stats.append(curr_mssg) - - @check_print_stats_step - def print_stats(self, fraction_done: float): - """helper to print out the stats dictionary. - - Args: - fraction_done (float): fraction of steps executed in training iterations - """ - # print a new header line if there is a new key added - # self.handle_header() - - # generate a new stats reporting message - self.handle_stats(fraction_done) - - def dump_stats(self): - """Dump stats locally to a json file""" - raise NotImplementedError - - -def cursorup(x: int): - """utility tool to move the cursor up on the terminal - - Args: - x (int): amount of lines to move cursor upward - """ - print(f"\r\033[{x}A", end="") diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index 735f2687d..ecacd7567 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -3,6 +3,8 @@ Generic Writer class """ +import datetime +import enum import os from abc import abstractmethod from typing import Dict @@ -10,66 +12,147 @@ from typing import Dict import imageio import numpy as np import torch -from omegaconf import DictConfig +from omegaconf import DictConfig, ListConfig from torch.utils.tensorboard import SummaryWriter from torchtyping import TensorType import mattport.utils.writer -from mattport.utils import comms from mattport.utils.decorators import check_main_thread, decorate_all to8b = lambda x: (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8) EVENT_WRITERS = [] +EVENT_STORAGE = [] +GLOBAL_BUFFER = {} -def write_image(name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None) -> None: - """Write image to all writers - Args: - name (str): data identifier - image (TensorType["H", "W", "C"]): rendered image to write - step (int): the time step to log - group (str): the group e.g., "Loss", "Accuracy", "Time" - prefix (str): the prefix e.g., "train-", "test-" - """ - assert comms.is_main_process(), "Writing out with process other than main" - for writer in EVENT_WRITERS: - writer.write_image(name, image, step, group, prefix) - - -def write_scalar(name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None: - """Write scalar to all writers +class EventType(enum.Enum): + """Possible Event types that can be logged Args: - name (str): data identifier - step (int): the time step to log - group (str): the group e.g., "Loss", "Accuracy", "Time" - prefix (str): the prefix e.g., "train-", "test-" + enum (_type_): _description_ """ - assert comms.is_main_process(), "Writing out with process other than main" - for writer in EVENT_WRITERS: - writer.write_scalar(name, scalar, step, group, prefix) + + IMAGE = "write_image" + SCALAR = "write_scalar" + DICT = "write_scalar_dict" -def write_scalar_dict(scalar_dict: Dict[str, float], step: int, group: str = None, prefix: str = None) -> None: - """Write dictionary of scalars to all writers +@check_main_thread +def put_image(name, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None): + """Setter function to place images into the queue to be written out""" + EVENT_STORAGE.append( + {"name": name, "write_type": EventType.IMAGE, "event": image, "step": step, "group": group, "prefix": prefix} + ) - Args: - scalar_dict (dict): dictionary containing all scalar values with key names and quantities - step (int): the time step to log - group (str): the group e.g., "Loss", "Accuracy", "Time" - prefix (str): the prefix e.g., "train-", "test-" + +@check_main_thread +def put_scalar(name: str, scalar: float, step: int, group: str = None, prefix: str = None): + """Setter function to place scalars into the queue to be written out""" + GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"] + GLOBAL_BUFFER["events"][name] = scalar + EVENT_STORAGE.append( + {"name": name, "write_type": EventType.SCALAR, "event": scalar, "step": step, "group": group, "prefix": prefix} + ) + + +@check_main_thread +def put_dict(name: str, scalar_dict: float, step: int, group: str = None, prefix: str = None): + """Setter function to place a dictionary of scalars into the queue to be written out""" + EVENT_STORAGE.append( + { + "name": name, + "write_type": EventType.DICT, + "event": scalar_dict, + "step": step, + "group": group, + "prefix": prefix, + } + ) + + +@check_main_thread +def put_time( + name: str, + start_time: float, + end_time: float, + step: int, + group: str = None, + prefix: str = None, + avg_over_iters: bool = False, + avg_over_batch: int = None, + update_eta: bool = False, +): + """Setter function to place a time element into the queue to be written out + + Processes the time info according to the options: + avg_over_iters (bool): if True, calculate and record a running average of the times + avg_over_batch (int): if set, the size of the batch for which we take the average over (batch/second) + update_eta (bool): if True, update the ETA. should only be set for the training iterations/s """ - assert comms.is_main_process(), "Writing out with process other than main" + GLOBAL_BUFFER["step"] = step + GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"] + val = end_time - start_time + if avg_over_batch: + val = avg_over_batch / val + + if avg_over_iters: + curr_event = GLOBAL_BUFFER["events"].get(name, {"buffer": [], "avg": 0}) + curr_buffer = curr_event["buffer"] + curr_avg = curr_event["avg"] + if len(curr_buffer) >= GLOBAL_BUFFER["max_history"]: + curr_buffer.pop(0) + curr_buffer.append(val) + curr_avg = sum(curr_buffer) / len(curr_buffer) + put_scalar(name, curr_avg, step, group, prefix) + GLOBAL_BUFFER["events"][name] = {"buffer": curr_buffer, "avg": curr_avg} + else: + GLOBAL_BUFFER["events"][name] = val + put_scalar(name, val, step, group, prefix) + + if update_eta: + ## NOTE: eta should be called with avg train iteration time + remain_iter = GLOBAL_BUFFER["max_iter"] - step + remain_time = remain_iter * GLOBAL_BUFFER["events"][name]["avg"] + GLOBAL_BUFFER["events"]["ETA"] = remain_time + put_scalar("ETA", remain_time, step, group, prefix) + + +@check_main_thread +def write_out_storage(): + """Function that writes all the events in storage to all the writer locations""" for writer in EVENT_WRITERS: - writer.write_scalar_dict(scalar_dict, step, group, prefix) + for event in EVENT_STORAGE: + write_func = getattr(writer, event["write_type"].value) + if event["write_type"] == EventType.DICT: + write_func(event["event"], event["step"], event["group"], event["prefix"]) + if isinstance(writer, LocalWriter): + continue + else: + write_func(event["name"], event["event"], event["step"], event["group"], event["prefix"]) + if isinstance(writer, LocalWriter): + break + EVENT_STORAGE.clear() +@check_main_thread def setup_event_writers(config: DictConfig) -> None: """Initialization of all event writers specified in config""" - if comms.is_main_process(): - for writer_type in config: - curr_writer = getattr(mattport.utils.writer, writer_type) - EVENT_WRITERS.append(curr_writer(config[writer_type].save_dir)) + logging_configs = config.logging.writer + for writer_type in logging_configs: + writer_class = getattr(mattport.utils.writer, writer_type) + curr_config = logging_configs[writer_type] + if writer_type == "LocalWriter": + curr_writer = writer_class(curr_config.save_dir, curr_config.stats_to_track, config) + else: + curr_writer = writer_class(curr_config.save_dir) + EVENT_WRITERS.append(curr_writer) + + ## configure all the global buffer basic information + GLOBAL_BUFFER["max_iter"] = config.graph.max_num_iterations + GLOBAL_BUFFER["max_history"] = config.logging.max_history + GLOBAL_BUFFER["steps_per_log"] = config.logging.steps_per_log + GLOBAL_BUFFER["new_key"] = True + GLOBAL_BUFFER["events"] = {} def get_tensorboard_name(name: str, group: str = None, prefix: str = None): @@ -165,16 +248,94 @@ class TensorboardWriter(Writer): self.tb_writer.add_scalar(tensorboard_name, scalar, step) +def _cursorup(x: int): + """utility tool to move the cursor up on the terminal + + Args: + x (int): amount of lines to move cursor upward + """ + print(f"\r\033[{x}A", end="\x1b[1K\r") + + +def _format_time(v): + if v < 1.0: + ms = v * 1e3 + v = f"{ms:0.4f} ms" + else: + v = str(datetime.timedelta(seconds=v)) + return v + + +def _consolidate_events(): + latest_map = {} + for event in EVENT_STORAGE: + name = event["name"] + latest_map[name] = event["event"] + return latest_map + + @decorate_all([check_main_thread]) class LocalWriter(Writer): """Local Writer Class""" + def __init__(self, save_dir: str, stats_to_track: ListConfig, config: DictConfig): + super().__init__(save_dir) + self.stats_to_track = stats_to_track + self.past_stats = [] + self.max_mssg_len = 0 + self.config = config + def write_image( self, name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None ) -> None: - image = to8b(image) - image_path = os.path.join(self.save_dir, f"{name}.jpg") - imageio.imwrite(image_path, np.uint8(image.cpu().numpy() * 255.0)) + if name in self.stats_to_track: + image = to8b(image) + image_path = os.path.join(self.save_dir, f"{name}.jpg") + imageio.imwrite(image_path, np.uint8(image.cpu().numpy() * 255.0)) def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None: - raise NotImplementedError + latest_map = _consolidate_events() + self.handle_header(latest_map) + self.handle_stats(latest_map) + GLOBAL_BUFFER["new_key"] = False + + def handle_header(self, latest_map): + """helper to handle the printing of the header labels""" + if len(self.past_stats) == 0 or GLOBAL_BUFFER["new_key"]: + mssg = f"{'Step (% Done)':<20}" + for name, _ in latest_map.items(): + if name in self.stats_to_track: + mssg += f"{name:<20} " + if len(self.past_stats) > 1: + _cursorup(len(self.past_stats) - 1) + print(mssg) + print("-" * len(mssg)) + if len(self.past_stats) > 0: + for mssg in self.past_stats: + print(mssg) + self.max_mssg_len = len(mssg) + + def handle_stats(self, latest_map, padding=" "): + """helper to print out the stats in a readable format""" + step = GLOBAL_BUFFER["step"] + fraction_done = step / GLOBAL_BUFFER["max_iter"] + curr_mssg = f"{step} ({fraction_done*100:.02f}%)" + curr_mssg = f"{curr_mssg:<20}" + for name, v in latest_map.items(): + if name in self.stats_to_track: + if "(time)" in name: + v = _format_time(v) + else: + v = f"{v:0.4f}" + curr_mssg += f"{v:<20} " + + # update the history buffer + if len(self.past_stats) > GLOBAL_BUFFER["max_history"]: + self.past_stats.pop(0) + _cursorup(len(self.past_stats)) + elif len(self.past_stats) > 1: + _cursorup(len(self.past_stats) - 1) + + for mssg in self.past_stats: + print(f"{mssg:{padding}<{self.max_mssg_len}}") + self.past_stats.append(curr_mssg)