diff --git a/.vscode/launch.json b/.vscode/launch.json index 6af3fad15..453912d23 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -51,6 +51,16 @@ "console": "integratedTerminal", "args": ["mipnerf_360"] }, + { + "name": "Python: run_train tensorf (gpu)", + "type": "python", + "request": "launch", + "program": "scripts/run_train.py", + "console": "integratedTerminal", + "args": [ + "tensorf" + ] + }, { "name": "Python: run_train nerfw (gpu)", "type": "python", diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 1b0678512..5012c4aaf 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -35,6 +35,7 @@ from nerfactory.dataloaders.datasets import ( from nerfactory.models.base import Model from nerfactory.models.instant_ngp import NGPModel from nerfactory.models.nerfw import NerfWModel +from nerfactory.models.tensorf import TensoRFModel from nerfactory.optimizers.schedulers import ExponentialDecaySchedule from nerfactory.pipelines.base import Pipeline from nerfactory.utils import writer @@ -293,6 +294,17 @@ class NerfWModelConfig(ModelConfig): uncertainty_min: float = 0.03 +@dataclass +class TensoRFModelConfig(ModelConfig): + """TensoRF model config""" + + _target: Type = TensoRFModel + init_resolution: int = 128 + final_resolution: int = 200 + upsampling_iters: Tuple[int, ...] = (5000, 5500, 7000) + loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss_coarse": 1.0, "feature_loss": 8e-5}) + + # Pipeline related configs @dataclass class PipelineConfig(InstantiateConfig): diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index 9b02c628e..09331f447 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -28,6 +28,8 @@ from nerfactory.configs.base import ( NerfWModelConfig, OptimizerConfig, PipelineConfig, + SchedulerConfig, + TensoRFModelConfig, TrainerConfig, VanillaDataloaderConfig, ) @@ -122,3 +124,28 @@ base_configs["vanilla_nerf"] = Config( model=ModelConfig(_target=NeRFModel), ), ) + +base_configs["tensorf"] = Config( + method_name="tensorf", + trainer=TrainerConfig(mixed_precision=True), + pipeline=PipelineConfig( + dataloader=VanillaDataloaderConfig( + train_dataset=BlenderDatasetConfig(), + ), + model=TensoRFModelConfig(), + ), + optimizers={ + "fields": { + "optimizer": OptimizerConfig(lr=0.001), + "scheduler": SchedulerConfig(lr_final=0.00005, max_steps=15000), + }, + "position_encoding": { + "optimizer": OptimizerConfig(lr=0.02), + "scheduler": SchedulerConfig(lr_final=0.005, max_steps=15000), + }, + "direction_encoding": { + "optimizer": OptimizerConfig(lr=0.02), + "scheduler": SchedulerConfig(lr_final=0.005, max_steps=15000), + }, + }, +) diff --git a/nerfactory/dataloaders/base.py b/nerfactory/dataloaders/base.py index b6c42e073..11bca9ee5 100644 --- a/nerfactory/dataloaders/base.py +++ b/nerfactory/dataloaders/base.py @@ -32,6 +32,7 @@ from nerfactory.dataloaders.image_dataset import ImageDataset, PanopticImageData from nerfactory.dataloaders.image_sampler import CacheImageSampler from nerfactory.dataloaders.pixel_sampler import PixelSampler from nerfactory.models.modules.ray_generator import RayGenerator +from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfactory.utils.misc import IterableWrapper @@ -177,6 +178,12 @@ class Dataloader(nn.Module): """ raise NotImplementedError + def get_training_callbacks( # pylint:disable=no-self-use + self, training_callback_attributes: TrainingCallbackAttributes # pylint: disable=unused-argument + ) -> List[TrainingCallback]: + """Returns a list of callbacks to be used during training.""" + return [] + @abstractmethod def get_param_groups(self) -> Dict[str, List[Parameter]]: # pylint: disable=no-self-use """Get the param groups for the dataloader. diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 958e7bf61..ff9c496f8 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -29,11 +29,14 @@ from torch.cuda.amp.grad_scaler import GradScaler from torch.nn.parallel import DistributedDataParallel as DDP from nerfactory.configs import base as cfg -from nerfactory.dataloaders.structs import DatasetInputs from nerfactory.optimizers.optimizers import Optimizers, setup_optimizers from nerfactory.pipelines.base import Pipeline from nerfactory.utils import profiler, writer -from nerfactory.utils.callbacks import Callback +from nerfactory.utils.callbacks import ( + TrainingCallback, + TrainingCallbackAttributes, + TrainingCallbackLocation, +) from nerfactory.utils.decorators import check_main_thread from nerfactory.utils.writer import EventName, TimeWriter from nerfactory.viewer.server import viewer_utils @@ -76,8 +79,6 @@ class Trainer: if self.device == "cpu": self.mixed_precision = False logging.warning("Mixed precision is disabled for CPU training.") - # dataset variables - self.dataset_inputs_train: DatasetInputs # model variables self.pipeline: Pipeline self.optimizers: Optimizers @@ -89,7 +90,7 @@ class Trainer: banner_messages = [f"Viewer at: {self.visualizer_state.viewer_url}"] self.grad_scaler = GradScaler(enabled=self.mixed_precision) # training callbacks - self.callbacks: List[Callback] + self.callbacks: List[TrainingCallback] # logging variables writer.setup_event_writers( config.logging, max_iter=config.trainer.max_num_iterations, banner_messages=banner_messages @@ -113,8 +114,9 @@ class Trainer: ) dist.barrier(device_ids=[self.local_rank]) - # TODO(ethan): do this for pipeline, not pipeline.model - self.callbacks = self.pipeline.model.get_training_callbacks() + self.callbacks = self.pipeline.get_training_callbacks( + TrainingCallbackAttributes(optimizers=self.optimizers, grad_scaler=self.grad_scaler, pipeline=self.pipeline) + ) def train(self) -> None: """Train the model.""" @@ -126,10 +128,17 @@ class Trainer: num_iterations = self.config.trainer.max_num_iterations for step in range(self.start_step, self.start_step + num_iterations): - # Note: if visualizer used, the rendering of the visualizer will be included in the iteration train time + # if the visualizer used, the rendering of the visualizer will be included in the iteration train time with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as t: + + # training callbacks before the training iteration + for callback in self.callbacks: + callback.run_callback_at_location( + step, location=TrainingCallbackLocation.BEFORE_TRAIN_ITERATION + ) + loss_metric_dict = self.train_iteration(step) - with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as t: + with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _: self.visualizer_state.update_scene(step, self.pipeline.model) train_num_rays_per_batch = self.config.pipeline.dataloader.train_num_rays_per_batch @@ -219,10 +228,7 @@ class Trainer: self.grad_scaler.scale(loss).backward() # type: ignore self.optimizers.optimizer_scaler_step_all(self.grad_scaler) self.grad_scaler.update() - self.optimizers.scheduler_step_all(step) - for callback in self.callbacks: - callback.after_step(step) # Merging loss and metrics dict into a single output. loss_dict["loss"] = loss diff --git a/nerfactory/fields/modules/encoding.py b/nerfactory/fields/modules/encoding.py index c195b9f75..6aad364b5 100644 --- a/nerfactory/fields/modules/encoding.py +++ b/nerfactory/fields/modules/encoding.py @@ -366,13 +366,21 @@ class TensorVMEncoding(Encoding): init_scale: Initialization scale. Defaults to 0.1. """ - def __init__(self, resolution: int = 256, num_components: int = 24, init_scale: float = 0.1) -> None: + plane_coef: TensorType[3, "num_components", "resolution", "resolution"] + line_coef: TensorType[3, "num_components", "resolution", 1] + + def __init__( + self, + resolution: int = 128, + num_components: int = 24, + init_scale: float = 0.1, + ) -> None: super().__init__(in_dim=3) self.resolution = resolution self.num_components = num_components - # TODO Learning rates should be different for these + # TODO(terrance): Learning rates should be different for these self.plane_coef = nn.Parameter(init_scale * torch.randn((3, num_components, resolution, resolution))) self.line_coef = nn.Parameter(init_scale * torch.randn((3, num_components, resolution, 1))) @@ -404,14 +412,12 @@ class TensorVMEncoding(Encoding): Args: resolution: Target resolution. """ - - self.plane_coef.data = F.interpolate( + plane_coef = F.interpolate( self.plane_coef.data, size=(resolution, resolution), mode="bilinear", align_corners=True ) - self.line_coef.data = F.interpolate( - self.line_coef.data, size=(resolution, 1), mode="bilinear", align_corners=True - ) - + line_coef = F.interpolate(self.line_coef.data, size=(resolution, 1), mode="bilinear", align_corners=True) + # TODO(ethan): are these torch.nn.Parameters needed? + self.plane_coef, self.line_coef = torch.nn.Parameter(plane_coef), torch.nn.Parameter(line_coef) self.resolution = resolution diff --git a/nerfactory/models/base.py b/nerfactory/models/base.py index bf34a7fec..a502738c8 100644 --- a/nerfactory/models/base.py +++ b/nerfactory/models/base.py @@ -30,7 +30,7 @@ from nerfactory.cameras.rays import RayBundle from nerfactory.configs import base as cfg from nerfactory.dataloaders.structs import SceneBounds from nerfactory.fields.density_fields.density_grid import DensityGrid -from nerfactory.utils.callbacks import Callback +from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfactory.utils.misc import get_masked_dict, is_not_none @@ -76,7 +76,9 @@ class Model(nn.Module): """Returns the device that the model is on.""" return self.device_indicator_param.device - def get_training_callbacks(self) -> List[Callback]: # pylint:disable=no-self-use + def get_training_callbacks( # pylint:disable=no-self-use + self, training_callback_attributes: TrainingCallbackAttributes # pylint: disable=unused-argument + ) -> List[TrainingCallback]: """Returns a list of callbacks that run functions at the specified training iterations.""" return [] @@ -92,7 +94,6 @@ class Model(nn.Module): update_every_num_iters=self.config.density_field_params["update_every_num_iters"], ) - @abstractmethod def populate_fields(self): """Set the fields.""" diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index afbbf93e2..d1617f056 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -35,7 +35,11 @@ from nerfactory.models.base import Model from nerfactory.models.modules.ray_sampler import NGPSpacedSampler from nerfactory.optimizers.loss import MSELoss from nerfactory.utils import colors, misc, visualization, writer -from nerfactory.utils.callbacks import Callback +from nerfactory.utils.callbacks import ( + TrainingCallback, + TrainingCallbackAttributes, + TrainingCallbackLocation, +) class NGPModel(Model): @@ -51,13 +55,16 @@ class NGPModel(Model): self.field = None super().__init__(config=config, **kwargs) - def get_training_callbacks(self) -> List[Callback]: + def get_training_callbacks( + self, training_callback_attributes: TrainingCallbackAttributes + ) -> List[TrainingCallback]: assert self.density_field is not None return [ - Callback( + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], update_every_num_iters=self.density_field.update_every_num_iters, func=self.density_field.update_density_grid, - density_eval_func=self.field.density_fn, # type: ignore + kwargs={"density_eval_func": self.field.density_fn}, # type: ignore ) ] diff --git a/nerfactory/models/tensorf.py b/nerfactory/models/tensorf.py index 044a152fd..98a36b1a1 100644 --- a/nerfactory/models/tensorf.py +++ b/nerfactory/models/tensorf.py @@ -14,5 +14,218 @@ """ TensorRF implementation. -TODO: """ + +from __future__ import annotations + +from typing import Dict, List + +import numpy as np +import torch +from torch.nn import Parameter +from torchmetrics import PeakSignalNoiseRatio +from torchmetrics.functional import structural_similarity_index_measure +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity + +from nerfactory.cameras.rays import RayBundle +from nerfactory.configs import base as cfg +from nerfactory.fields.modules.encoding import TensorVMEncoding +from nerfactory.fields.modules.field_heads import FieldHeadNames +from nerfactory.fields.nerf_field import NeRFField +from nerfactory.models.base import Model +from nerfactory.models.modules.ray_sampler import PDFSampler, UniformSampler +from nerfactory.models.modules.scene_colliders import NearFarCollider +from nerfactory.optimizers.loss import L1Loss, MSELoss +from nerfactory.optimizers.optimizers import Optimizers +from nerfactory.renderers.renderers import ( + AccumulationRenderer, + DepthRenderer, + RGBRenderer, +) +from nerfactory.utils import colors, misc, visualization, writer +from nerfactory.utils.callbacks import ( + TrainingCallback, + TrainingCallbackAttributes, + TrainingCallbackLocation, +) + + +class TensoRFModel(Model): + """ + TensoRF Model + """ + + def __init__( + self, + config: cfg.TensoRFModelConfig, + **kwargs, + ) -> None: + self.init_resolution = config.init_resolution + self.upsampling_iters = config.upsampling_iters + self.upsampling_steps = ( + np.round( + np.exp( + np.linspace( + np.log(config.init_resolution), + np.log(config.final_resolution), + len(config.upsampling_iters) + 1, + ) + ) + ) + .astype("int") + .tolist()[1:] + ) + super().__init__(config=config, **kwargs) + + def get_training_callbacks( + self, training_callback_attributes: TrainingCallbackAttributes + ) -> List[TrainingCallback]: + + # the callback that we want to run every X iterations after the training iteration + def reinitialize_optimizer( + self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument + ): + resolution = self.upsampling_steps.pop(0) + + # upsample the position and direction grids + # TODO(ethan): ask Brent how to get typing to work on this... the Encoding base class type + # in NeRFField is causing the issue + self.field.position_encoding.upsample_grid(resolution) + self.field.direction_encoding.upsample_grid(resolution) + + # reinitialize the optimizer + optimizers_config = training_callback_attributes.optimizers.config + training_callback_attributes.optimizers = Optimizers( + optimizers_config, training_callback_attributes.pipeline.get_param_groups() + ) + # TODO(ethan): do something with the learning rate + # we don't want to reinitialize the learning rate each time + + callbacks = [ + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + iters=self.upsampling_iters, + func=reinitialize_optimizer, + args=[self, training_callback_attributes], + ) + ] + return callbacks + + def populate_misc_modules(self): + # fields + position_encoding = TensorVMEncoding( + resolution=self.init_resolution, + num_components=96, + ) + direction_encoding = TensorVMEncoding( + resolution=self.init_resolution, + num_components=96, + ) + + self.field = NeRFField( + position_encoding=position_encoding, + direction_encoding=direction_encoding, + base_mlp_num_layers=2, + base_mlp_layer_width=128, + ) + + # samplers + self.sampler_uniform = UniformSampler( + num_samples=self.config.num_coarse_samples, density_field=self.density_field + ) + self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples, density_field=self.density_field) + + # renderers + self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) + self.renderer_accumulation = AccumulationRenderer() + self.renderer_depth = DepthRenderer() + + # losses + self.rgb_loss = MSELoss() + self.feature_loss = L1Loss() + + # metrics + self.psnr = PeakSignalNoiseRatio(data_range=1.0) + self.ssim = structural_similarity_index_measure + self.lpips = LearnedPerceptualImagePatchSimilarity() + + # colliders + if self.config.enable_collider: + self.collider = NearFarCollider( + near_plane=self.config.collider_params["near_plane"], far_plane=self.config.collider_params["far_plane"] + ) + + def get_param_groups(self) -> Dict[str, List[Parameter]]: + param_groups = {} + param_groups["fields"] = list(self.field.parameters()) + param_groups["position_encoding"] = list(self.field.position_encoding.parameters()) + param_groups["direction_encoding"] = list(self.field.direction_encoding.parameters()) + return param_groups + + def get_outputs(self, ray_bundle: RayBundle): + + # uniform sampling + ray_samples_uniform = self.sampler_uniform(ray_bundle) + + # coarse field: + field_outputs_coarse = self.field.forward(ray_samples_uniform) + weights = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) + rgb = self.renderer_rgb( + rgb=field_outputs_coarse[FieldHeadNames.RGB], + weights=weights, + ) + accumulation = self.renderer_accumulation(weights) + depth = self.renderer_depth(weights, ray_samples_uniform) + + outputs = {"rgb": rgb, "accumulation": accumulation, "depth": depth} + return outputs + + def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients) -> Dict[str, torch.Tensor]: + # Scaling metrics by coefficients to create the losses. + device = outputs["rgb"].device + image = batch["image"].to(device) + assert isinstance(self.field.position_encoding, TensorVMEncoding) + + rgb_loss = self.rgb_loss(image, outputs["rgb"]) + plane_coef = self.field.position_encoding.plane_coef + line_coef = self.field.position_encoding.line_coef + + plane_feature_loss = self.feature_loss(plane_coef, torch.zeros_like(plane_coef)) + line_feature_loss = self.feature_loss(line_coef, torch.zeros_like(line_coef)) + + loss_dict = {"rgb_loss": rgb_loss, "feature_loss": plane_feature_loss + line_feature_loss} + loss_dict = misc.scale_dict(loss_dict, loss_coefficients) + return loss_dict + + def log_test_image_outputs(self, image_idx, step, batch, outputs): + image = batch["image"] + rgb = outputs["rgb"] + acc = visualization.apply_colormap(outputs["accumulation"]) + depth = visualization.apply_depth_colormap( + outputs["depth"], + accumulation=outputs["accumulation"], + near_plane=self.config.collider_params["near_plane"], + far_plane=self.config.collider_params["far_plane"], + ) + + combined_rgb = torch.cat([image, rgb], dim=1) + + writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step) + writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=acc, step=step) + writer.put_image(name=f"depth/image_idx_{image_idx}", image=depth, step=step) + + # Switch images from [H, W, C] to [1, C, H, W] for metrics computations + image = torch.moveaxis(image, -1, 0)[None, ...] + rgb = torch.moveaxis(rgb, -1, 0)[None, ...] + + psnr = self.psnr(image, rgb) + ssim = self.ssim(image, rgb) + lpips = self.lpips(image, rgb) + + writer.put_scalar(name=f"psnr/val_{image_idx}", scalar=float(psnr), step=step) + writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(ssim), step=step) # type: ignore + writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(lpips), step=step) + + writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(psnr), step=step) + + return psnr.item() diff --git a/nerfactory/models/vanilla_nerf.py b/nerfactory/models/vanilla_nerf.py index 332630ec9..70ba6ce46 100644 --- a/nerfactory/models/vanilla_nerf.py +++ b/nerfactory/models/vanilla_nerf.py @@ -41,7 +41,11 @@ from nerfactory.renderers.renderers import ( RGBRenderer, ) from nerfactory.utils import colors, misc, visualization, writer -from nerfactory.utils.callbacks import Callback +from nerfactory.utils.callbacks import ( + TrainingCallback, + TrainingCallbackAttributes, + TrainingCallbackLocation, +) class NeRFModel(Model): @@ -69,20 +73,22 @@ class NeRFModel(Model): **kwargs, ) - def get_training_callbacks(self) -> List[Callback]: + def get_training_callbacks( + self, training_callback_attributes: TrainingCallbackAttributes # pylint: disable=unused-argument + ) -> List[TrainingCallback]: if self.field_coarse is None: raise ValueError("populate fields must be called before get_training_callbacks.") - callbacks = [] if self.density_field is not None: callbacks = [ - Callback( + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], update_every_num_iters=self.density_field.update_every_num_iters, func=self.density_field.update_density_grid, - density_eval_func=self.field_coarse.density_fn, + kwargs={"density_eval_func": self.field.density_fn}, # type: ignore ) ] - return callbacks # type: ignore + return callbacks def populate_fields(self): """Set the fields.""" diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 8d44e4789..774e24d65 100644 --- a/nerfactory/pipelines/base.py +++ b/nerfactory/pipelines/base.py @@ -27,6 +27,7 @@ from nerfactory.configs import base as cfg from nerfactory.dataloaders.base import Dataloader from nerfactory.models.base import Model from nerfactory.utils import profiler +from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes class Pipeline(nn.Module): @@ -116,6 +117,15 @@ class Pipeline(nn.Module): state = {key.replace("module.", ""): value for key, value in loaded_state.items()} self.load_state_dict(state) # type: ignore + def get_training_callbacks( + self, training_callback_attributes: TrainingCallbackAttributes + ) -> List[TrainingCallback]: + """Returns the training callbacks from both the Dataloader and the Model.""" + dataloader_callbacks = self.dataloader.get_training_callbacks(training_callback_attributes) + model_callbacks = self.model.get_training_callbacks(training_callback_attributes) + callbacks = dataloader_callbacks + model_callbacks + return callbacks + def get_param_groups(self) -> Dict[str, List[Parameter]]: """Get the param groups for the pipeline. diff --git a/nerfactory/utils/callbacks.py b/nerfactory/utils/callbacks.py index f48d8d84b..874c6aa17 100644 --- a/nerfactory/utils/callbacks.py +++ b/nerfactory/utils/callbacks.py @@ -16,13 +16,36 @@ Callback code used for training iterations """ -from typing import Callable +from dataclasses import InitVar, dataclass +from enum import Enum, auto +from inspect import signature +from typing import Callable, Dict, List, Optional, Tuple -class Callback: +@dataclass +class TrainingCallbackAttributes: + """Attributes that can be used to configure training callbacks. + The callbacks can be specified in the Dataloader or Model implementations. + Instead of providing access to the entire Trainer object, we only provide these attributes. + This should be least prone to errors and fairly clean from a user perspective.""" + + # TODO(ethan): type this without circular imports + optimizers: Optional[InitVar] + grad_scaler: Optional[InitVar] + pipeline: Optional[InitVar] + + +class TrainingCallbackLocation(Enum): + """Enum for specifying where the training callback should be run.""" + + BEFORE_TRAIN_ITERATION = auto() + AFTER_TRAIN_ITERATION = auto() + + +class TrainingCallback: """Callback class used during training. The function 'func' with 'args' and 'kwargs' will be called every 'update_every_num_iters' training iterations, - including at iteration 0. + including at iteration 0. The function is called after the training iteration. Args: update_every_num_iters: How often to call the function `func`. @@ -31,26 +54,39 @@ class Callback: kwargs: kwargs for the function 'func'. """ - def __init__(self, update_every_num_iters: int, func: Callable, *args, **kwargs): - # TODO(ethan): how do we type args and kwargs? + def __init__( + self, + where_to_run: List[TrainingCallbackLocation], + func: Callable, + update_every_num_iters: Optional[int] = None, + iters: Optional[Tuple[int, ...]] = None, + args: Optional[List] = None, + kwargs: Optional[Dict] = None, + ): + assert ( + "step" in signature(func).parameters.keys() + ), f"'step: int' must be an argument in the callback function 'func': {func.__name__}" + self.where_to_run = where_to_run self.update_every_num_iters = update_every_num_iters + self.iters = iters self.func = func - self.args = args - self.kwargs = kwargs + self.args = args if args is not None else [] + self.kwargs = kwargs if kwargs is not None else {} - def before_step(self, step: int): - """callback to run before training step + def run_callback(self, step: int): + """Callback to run after training step Args: step (int): current iteration step """ - raise NotImplementedError + if self.update_every_num_iters is not None: + if step % self.update_every_num_iters == 0: + self.func(*self.args, **self.kwargs, step=step) + elif self.iters is not None: + if step in self.iters: + self.func(*self.args, **self.kwargs, step=step) - def after_step(self, step: int): - """callback to run after training step - - Args: - step (int): current iteration step - """ - if step % self.update_every_num_iters == 0: - self.func(*self.args, **self.kwargs, step=step) + def run_callback_at_location(self, step: int, location: TrainingCallbackLocation): + """Runs the callback if it's supposed to be run at the given location.""" + if location in self.where_to_run: + self.run_callback(step=step) diff --git a/tests/fields/modules/test_encodings.py b/tests/fields/modules/test_encodings.py index c9f8e0e5d..15d5526cc 100644 --- a/tests/fields/modules/test_encodings.py +++ b/tests/fields/modules/test_encodings.py @@ -95,6 +95,7 @@ def test_tensor_vm_encoder(): assert encoded.shape == (6, 3, out_dim) encoder.upsample_grid(resolution=64) + encoder.upsample_grid(resolution=128) def test_tensor_cp_encoder():