task.AtomisticTask

class task.AtomisticTask(*args: Any, **kwargs: Any)

The basic learning task in SchNetPack, which ties model, loss and optimizer together.

Parameters:

• model – the neural network model

• outputs – list of outputs an optional loss functions

• optimizer_cls – type of torch optimizer,e.g. torch.optim.Adam

• optimizer_args – dict of optimizer keyword arguments

• scheduler_cls – type of torch learning rate scheduler

• scheduler_args – dict of scheduler keyword arguments

• scheduler_monitor – name of metric to be observed for ReduceLROnPlateau

• warmup_steps – number of steps used to increase the learning rate from zero linearly to the target learning rate at the beginning of training