Trainers

`TrainingConfig(BATCH_SIZE=2 13, TEST_BATCH_SIZE=2 15, N_EPOCHS=100, LEARNING_RATE=0.001, DEVICE='cpu', OPTIMIZER='Adam', MATMUL_PRECISION='high', COMPILE=False)` `dataclass`

Configuration object storing hyperparameters for a Trainer class.

Attributes:

Name	Type	Description
`BATCH_SIZE`	`int`	Number of samples per batch in the training data loader. Defaults to 2**13 (= 8192).
`TEST_BATCH_SIZE`	`int`	Number of samples per batch in the test data loader. Defaults to 2**15 (= 32768).
`N_EPOCHS`	`int`	Number of passes over the training dataset. Defaults to 100.
`LEARNING_RATE`	`float`	Learning rate given to the optimizer. Defaults to 1e-3.
`DEVICE`	`str`	Torch device on which training runs ("cpu", "cuda", ...). A warning is emitted if set to "cpu" while a GPU is available. Defaults to "cpu".
`OPTIMIZER`	`str`	Name of the optimizer (case-insensitive). One of "adam", "sgd", "muon" or "custom". If "custom" is selected, the user must set `trainer.optimizer` before calling the `train` method). Unknown names fall back to Adam. Defaults to "Adam".
`MATMUL_PRECISION`	`str`	Passed to `torch.set_float32_matmul_precision` at the start of training. Choices are ["highest", "high", "medium"]. "high" enables TF32 matrix products on recent GPUs, which is faster but less precise. Defaults to "high".
`COMPILE`	`bool`	Whether to wrap the model with `torch.compile` for training. Compiling the model will greatly improve performance but will prevent the use of double differentiation, i.e. differential losses (like the eikonal loss). Defaults to False.

`Trainer(config)`

Base trainer class. Implements the optimization loop for a given neural implicit model.

This class is not meant to be used directly, but to be inherited from. Subclasses need to implement the abstract methods forward_train_batch and forward_test_batch.

Attributes:

Name	Type	Description
`config`	`TrainingConfig`	The training hyperparameters.
`train_data_loader`	`DataLoader`	Loader over the training data. Set by `set_training_data`.
`test_data_loader`	`DataLoader`	Loader over the test data. Set by `set_test_data`.
`optimizer`	`Optimizer`	The optimizer. Built at the start of `train`, unless `OPTIMIZER` is "custom", in which case it must be assigned by the user.
`scheduler`	`LRScheduler`	The learning rate scheduler, if any (see `add_scheduler`).
`callbacks`	`list[Callback]`	Callbacks called at different steps of the training loop.

Parameters:

Name	Type	Description	Default
`config`	`TrainingConfig`	The trainer parameters as a configuration object. If None, the default `TrainingConfig()` is used.	required

`add_callbacks(*args)`

Adds callbacks to the trainer. All added callbacks are invoked at the beginning and the end of each training epoch.

Parameters:

Name	Type	Description	Default
`*args`	`Callback`	The callbacks, given either as several arguments or as a single iterable.	`()`

Raises:

Type	Description
`AssertionError`	if one of the arguments is not a `Callback`.

`add_scheduler(scheduler_cls, *args, **kwargs)`

Registers a learning rate scheduler. It is instantiated at the start of train as scheduler_cls(optimizer, *args, **kwargs) and stepped once at the end of each epoch.

Parameters:

Name	Type	Description	Default
`scheduler_cls`	`type`	The scheduler class (not an instance), for instance `torch.optim.lr_scheduler.StepLR`.	required
`*args`		Positional arguments given to the scheduler after the optimizer.	`()`
`**kwargs`		Keyword arguments given to the scheduler.	`{}`

Example

trainer.add_scheduler(torch.optim.lr_scheduler.StepLR, step_size=10, gamma=0.5)

`evaluate_model(model)`

Evaluates the model on the test dataset. Does nothing if no test data was provided.

Sums the losses returned by the forward_test_batch method over all test batches, then calls the callOnEndTest method of the callbacks.

Parameters:

Name	Type	Description	Default
`model`	`Module`	The model to evaluate.	required

`forward_test_batch(data, model)` `abstractmethod`

Computes the test loss over one batch. To be implemented by subclasses.

Parameters:

Name	Type	Description	Default
`data`	`list[Tensor]`	The tensors of one batch of the test dataset, already on `config.DEVICE`.	required
`model`	`Module`	The model being evaluated.	required

Returns:

Type	Description
	torch.Tensor: The scalar loss of the batch.

`forward_train_batch(data, model)` `abstractmethod`

Computes the training loss over one batch. To be implemented by subclasses.

Parameters:

Name	Type	Description	Default
`data`	`list[Tensor]`	The tensors of one batch of the training dataset, already on `config.DEVICE`.	required
`model`	`Module`	The model being trained.	required

Returns:

Type	Description
	torch.Tensor: The scalar loss of the batch, on which `backward()` is called.

`get_optimizer(model)`

Builds the optimizer described by config.OPTIMIZER, with learning rate config.LEARNING_RATE:

"adam": torch.optim.Adam. Also used for unknown names.
"sgd": torch.optim.SGD with momentum 0.9.
"muon": Muon on weights of dimension >= 2, and Adam (lr 5e-4) on biases and gains. Weight decay 0.01.
"custom": no optimizer is built.

Parameters:

Name	Type	Description	Default
`model`	`Module`	The model whose parameters are optimized.	required

Returns:

Type	Description
	torch.optim.Optimizer: The optimizer, or None if `config.OPTIMIZER` is "custom".

`set_test_data(data)`

Sets the test dataset, used by evaluate_model at the end of each epoch. Builds a data loader using config.TEST_BATCH_SIZE and config.NUM_DATALOADER_WORKERS.

Parameters:

Name	Type	Description	Default
`data`	`TensorDataset`	The test data.	required

Raises:

Type	Description
`Exception`	if `data` is not a `torch.utils.data.TensorDataset`.

`set_training_data(data, shuffle=True)`

Sets the training dataset. Builds a data loader using config.BATCH_SIZE and config.NUM_DATALOADER_WORKERS.

Parameters:

Name	Type	Description	Default
`data`	`TensorDataset`	The training data.	required
`shuffle`	`bool`	Whether to shuffle the data at every epoch. Defaults to True.	`True`

Raises:

Type	Description
`Exception`	if `data` is not a `torch.utils.data.TensorDataset`.

`step_scheduler()`

Steps the learning rate scheduler, if any. Prints the new learning rate when it changes.

`train(model, starting_epoch=0)`

Runs the optimization loop for config.N_EPOCHS epochs.

If some were added, callbacks are invoked during training:

callOnBeginTrain before the first epoch;
for each epoch: callOnBeginEpoch, callOnEndForward after each batch, callOnEndEpoch, then the scheduler is stepped and the model is evaluated on test data (callOnEndTest);
callOnEndTrain after the last epoch.

Parameters:

Name	Type	Description	Default
`model`	`Module`	The model to train. It should already be on `config.DEVICE`.	required
`starting_epoch`	`int`	Offset of the epoch counter. Used to resume a previous training. Defaults to 0.	`0`

Raises:

Type	Description
`Exception`	if no training data was provided, or if `config.OPTIMIZER` is "custom" and `optimizer` was not set.

`RegressionEikonalTrainer(config, eikonal_weight=0.01)`

Bases: Trainer

A simple trainer that combines mean square error regression to distance values with an eikonal loss.

The model is evaluated on a box [-1, 1]^d

Parameters:

Name	Type	Description	Default
`config`	`TrainingConfig`	The training hyperparameters.	required
`eikonal_weight`	`float`	weigth associated with the eikonal loss. Defaults to 0.01.	`0.01`

`SimpleRegressionTrainer(config)`

Bases: Trainer

A simple trainer that performs mean square error regression between the output of a model and some ground truth values.

`ImplicitSurfaceTrainer(config)`

Bases: Trainer

Training a neural implicit to approximate a shape as its zero level set. The loss functions and hyperparameters are adapted from [1] and [2]. Given a neural network $f$, The method optimizes for four loss functions:

\[\mathcal{L}_{\text{on}}[f] = \int_M |f(x)| dx\]

\[\mathcal{L}_{\text{normal}}[f] = \int_M ||\nabla f(x) - n(x)||\;dx\]

\[\mathcal{L}_{\text{eikonal}}[f] = \int_D (|| \nabla f(x)|| - 1)^2\;dx\]

\[\mathbb{L}_{\text{out}}[f] = \int_D \exp(-\rho\,|f(x)|)\;dx\]

This trainer assumes that the dataset is made of (points, normals) pairs which are sampled on the shape, for the evaluation of $$\mathcal{L}_{\text{on}}$ and $\mathcal{L}_{\text{normal}}$. The two other loss functions are evaluated for every batch at random positions in a [-1.5, 1.5]^d box.

Note

The forward_test_batch method is not implemented and does nothing.

References

[1] Implicit Neural Representations with Periodic Activation Functions, Sitzmann et al., 2020
[2] Implicit Geometric Regularization for Learning Shapes, Gropp et al., 2020

Attributes:

Name	Type	Description
`rho`	`float`	value of rho in the outside loss. Defaults to 100.
`weights`	`dict`	dictionnary of weights that are associated with each loss function. Keys and default values are: - "on" : 100. - "out" : 5. - "normal": 1. - "eikonal" : 1.

Parameters:

Name	Type	Description	Default
`config`	`TrainingConfig`	The training hyperparameters.	required

`hKRTrainer(config, margin=0.01, lmbd=100.0, test_mode='sdf')`

Bases: Trainer

Trainer minimizing the hinge-Kantorovitch-Rubinstein loss. For a Lipschitz model $f$ and some binary labels $y(x) \in \{-1, 1\}$, the hKR loss is defined as:

\[\mathcal{L}_{hKR}(x) = -yf(x) + \lambda \max(0, m - yf(x)).\]

Warning

This trainer only yields valid result when used on a Lipschitz architecture.

References

[1] Achieving robustness in classification using optimal transport with hinge regularization, Serrurier et al, 2021
[2] 1-Lipschitz Neural Distance Fields, Coiffier and Béthune, 2024

Parameters:

Name	Type	Description	Default
`config`	`TrainingConfig`	training configuration hyperparameters.	required
`margin`	`float`	margin (m) parameter for the hKR loss. Defaults to 1e-2.	`0.01`
`lmbd`	`float`	lambda parameter of the hKR loss. Defaults to 100..	`100.0`
`test_mode`	`str`	which loss to compute for test batches. Choices are [none, sdf, hkr]. Defaults to "sdf".	`'sdf'`

`SphereInitializeTrainer(n_batches, config, **kwargs)`

Bases: Trainer

Trainer used as an initialization step. Trains the neural network to approximate the SDF of a nD sphere

References

[1] SAL: Sign Agnostic Learning of Shapes From Raw Data, Atzmon and Lipman, 2020

Parameters:

Name	Type	Description	Default
`n_batches`	`int`	number of training batches (and optimizer step) to perform.	required
`config`	`TrainingConfig`	training hyperparameters.	required

Other Parameters:

Name	Type	Description
`radius`	`float`	radius of the sphere to consider. Defaults to 1.

Trainers

TrainingConfig(BATCH_SIZE=2 ** 13, TEST_BATCH_SIZE=2 ** 15, N_EPOCHS=100, LEARNING_RATE=0.001, DEVICE='cpu', OPTIMIZER='Adam', MATMUL_PRECISION='high', COMPILE=False) dataclass

Trainer(config)

add_callbacks(*args)

add_scheduler(scheduler_cls, *args, **kwargs)

evaluate_model(model)

forward_test_batch(data, model) abstractmethod

forward_train_batch(data, model) abstractmethod

get_optimizer(model)

set_test_data(data)

set_training_data(data, shuffle=True)

step_scheduler()

train(model, starting_epoch=0)

RegressionEikonalTrainer(config, eikonal_weight=0.01)

SimpleRegressionTrainer(config)

ImplicitSurfaceTrainer(config)

hKRTrainer(config, margin=0.01, lmbd=100.0, test_mode='sdf')

SphereInitializeTrainer(n_batches, config, **kwargs)

`TrainingConfig(BATCH_SIZE=2 13, TEST_BATCH_SIZE=2 15, N_EPOCHS=100, LEARNING_RATE=0.001, DEVICE='cpu', OPTIMIZER='Adam', MATMUL_PRECISION='high', COMPILE=False)` `dataclass`

`Trainer(config)`

`add_callbacks(*args)`

`add_scheduler(scheduler_cls, *args, **kwargs)`

`evaluate_model(model)`

`forward_test_batch(data, model)` `abstractmethod`

`forward_train_batch(data, model)` `abstractmethod`

`get_optimizer(model)`

`set_test_data(data)`

`set_training_data(data, shuffle=True)`

`step_scheduler()`

`train(model, starting_epoch=0)`

`RegressionEikonalTrainer(config, eikonal_weight=0.01)`

`SimpleRegressionTrainer(config)`

`ImplicitSurfaceTrainer(config)`

`hKRTrainer(config, margin=0.01, lmbd=100.0, test_mode='sdf')`

`SphereInitializeTrainer(n_batches, config, **kwargs)`