deepsensor.train.train

class Trainer(model, lr=5e-05)

Bases: object

Class for training ConvNP models with an Adam optimiser.

Parameters:

lr (float) – Learning rate

__call__(tasks, batch_size=None, progress_bar=False, tqdm_notebook=False)

Train model for one epoch.

set_gpu_default_device()

Set default GPU device for the backend.

Raises:

• RuntimeError – If no GPU is available.

• RuntimeError – If backend is not supported.

• NotImplementedError – If backend is not supported.

Returns:

None.

train_epoch(model, tasks, lr=5e-05, batch_size=None, opt=None, progress_bar=False, tqdm_notebook=False)

Train model for one epoch.

Parameters:

• model (ConvNP) – Model to train.

• tasks (List[Task]) – List of tasks to train on.

• lr (float, optional) – Learning rate, by default 5e-5.

• batch_size (int, optional) – Batch size. Defaults to None. If None, no batching is performed.

• opt (Optimizer, optional) – TF or Torch optimizer. Defaults to None. If None, tensorflow:tensorflow.keras.optimizer.Adam is used.

• progress_bar (bool, optional) – Whether to display a progress bar. Defaults to False.

• tqdm_notebook (bool, optional) – Whether to use a notebook progress bar. Defaults to False.

Returns:

List[float] – List of losses for each task/batch.