Engines

import pyblaze.nn as xnn

The engine module is PyBlaze’s most central component as it provides the engines driving the training and evaluation loops. The base class enables implementing custom engines. However, the available engines are already usable for a variety of use cases.

Base Class

class pyblaze.nn.Engine(model)

A base class for training and evaluating models as well as making predictions. Generally, this class should be seen as _binding_ between a model and some data. Models should always be wrapped in an engine, both when being trained, and when performing inference/evaluation. The engine ensures that the model’s environment is correct and prevents plenty of pitfalls.

A concrete implementation of this class is tailored to a specific type of data (e.g. independent, identically distributed data samples) and/or model types (e.g. GANs).

supports_multiple_gpus()

Returns whether the engine allows multiple GPUs to be used during training. By default, it returns True.

Returns

Whether multiple GPUs are allowed.

Return type

bool

__init__(model)

Initializes a new engine for a specified model.

Parameters

model (torch.nn.Module) – The model to train or evaluate.

train(train_data, val_data=None, epochs=20, val_iterations=None, eval_every=None, eval_train=False, eval_val=True, callbacks=None, metrics=None, gpu='auto', **kwargs)

Method for training the model with the supplied parameters.

Parameters

• train_data (torch.utils.data.DataLoader) – A data loader to obtain training data from. The samples yielded by the data loader depend on a specific engine implementation.

• val_data (torch.utisl.data..DataLoader, default: None) – A data loader to use for validation. If the loader is an infinite data loader, val_iterations must also be given. If not supplied, no validation will be performed.

• epochs (int, default: 20) – The number of epochs to train for. If the given data is an infinite data loader, this value defines the number of iterations (and should most probably be increased).

• val_iterations (int, default: None) – The number of iterations to perform for validation. Must be given only if val_data is an infinite data loader. Otherwise, the parameter is ignored.

• eval_every (int, default: None) – How many iterations should pass until validation is called again. If this is set when train_data is an iterable dataset, then it defines the number of iterations per epoch (i.e. before performing validation again). Otherwise, it defines the number of epochs to train before calling validation.

• eval_train (bool, default: False) – Whether to compute metrics (apart from the loss) for both the validation and the train data. If this flag is set to False, metrics will only be computed for the validation data.

• eval_val (bool, default: True) – Whether to compute the loss (apart from the metrics) for the validation data. If this flag is set to False, no validation loss will be computed.

• callbacks (list of (pyblaze.nn.TrainingCallback or pyblaze.nn.PredictionCallback),) – default: [] The callbacks to use for training and inference. The training callbacks and prediction callbacks will be filtered automatically.

• metrics (dict of str -> func, default: {}) – Metrics to compute during evaluation for the validation data (and potentially for the training data). The keys for the metrics define their name.

• gpu (str or bool or int or list of int, default: 'auto') – Governs, whether training and evaluation should be performed on a GPU. If set to True, the first GPU is selected. If set to an integer, the GPU with the specified index is used. If set to a list of integers, the specified GPUs are used to train and evaluate the model und multiple GPUs simultaneously. In this case, the batch sizes of the data loaders should be adjusted accordingly. If set to a string, the only valid value is ‘auto’. In this case, all available GPUs are used.

• kwargs (keyword arguments) – Additional keyword arguments dependent on the specific subclass. If prefixed with ‘eval_’, it will be passed to eval_batch() (without the prefix), otherwise to train_batch(). All keyword arguments may also be given as ParameterScheduler to enable dynamic changes in parameters. Note that the scheduler can also be used to alter values for ‘eval_’ keywords. However, they do so following the same schedule as the training keywords. In any case, consult the engine’s documentation if the parameter allows dynamic behavior. Semantically, it might not make sense to do.

Returns

A history object summarizing stats from the training.

Return type

pyblaze.nn.History

evaluate(data, iterations=None, metrics=None, callbacks=None, gpu='auto', **kwargs)

Evaluates the model on the given data and computes the supplied metrics.

Parameters

• data (torch.DataLoader) – A data loader to obtain evaluation samples from. The expected samples depend on a specific engine subclass.

• iterations (int, default: None) – The number of samples used for evaluating if the given data is an infinite data loader.

• metrics (dict of str -> func, default: {}) – The metrics to evaluate the model for. The keys define the names of the metrics when retrieving the evaluated result from the return parameter.

• callbacks (list of pyblaze.nn.PredictionCallback, default: []) – Callbacks to use while computing predictions. Usually, they are used for logging.

• gpu (str or bool or int or list of int, default: False) – Governs, whether training and evaluation should be performed on a GPU. If set to True, the GPU with the most amount of free memory is selected (if there are multiple GPUs). If set to an integer, the GPU with the specified index is used. If set to a list of integers, the specified GPUs are used to train and evaluate the model und multiple GPUs simultaneously. In this case, the batch sizes of the data loaders should be adjusted accordingly. If set to ‘auto’, all available GPUs are used. In case of None, the model will not be moved. Only use this option if you know what you are doing.

• kwargs (keyword arguments) – Additional arguments, passed directly to the eval_batch method.

Returns

An evaluation object, yielding as properties the metrics with their specified names.

Return type

pyblaze.nn.training.wrappers.Evaluation

predict(data, iterations=None, callbacks=None, gpu='auto', **kwargs)

Computes predictions for the given samples.

Parameters

• data (torch.DataLoader) – The data loader from which to obtain items for prediction. Note that in most cases it is desirable that the data loader does not shuffle items as predictions may then be out-of-order.

• iterations (int, default: None) – The (maximum) number of samples used for evaluating if the given data is an iterable dataset.

• callbacks (list of pyblaze.nn.PredictionCallback) – Callbacks which are called as prediction progresses.

• gpu (str or bool or int or list of int, default: 'auto') – Whether to use a (specific) GPU or multiple GPUs. If multiple GPUs are used, one process per GPU is started to minimize synchronization overhead. Make sure that using multiple GPUs makes up for this overhead. If False is specified, all cores of the computer are used to make predictions in parallel. In the case of ‘auto’, all available GPUs are used (if any).

• kwargs (keyword arguments) – Additional arguments passed directly to the predict_batch function.

Returns

The predictions made by the model as collated by collate_predictions.

Return type

object

abstract train_batch(data, **kwargs)

Runs a single step in training. If the training data represents an infinite dataset, this equals a single iteration, otherwise a mini-batch.

Parameters

data (object) – The data for the current iteration/mini-batch.

Returns

The loss computed for the batch. If the returned value is not float, overwrite collate_losses(). If you return a dict here, all keys prefixed with an underscore will not be passed to any callbacks but only be available in the collate_losses() method.

Return type

object

abstract eval_batch(data, **kwargs)

Runs a single step for inference. The data is either a mini-batch or a single iteration, depending on the data used for evaluation.

Parameters

• data (object) – The data for the current iteration/mini-batch.

• kwargs (keyword arguments) – Additional arguments dependent on the subclass and passed directly from the evaluation method.

Returns

The model output and possibly some correct target. If this is not a torch.Tensor, the collate_evals().

Return type

object

predict_batch(data, **kwargs)

Processes a single batch of inputs to make model predictions. The default implementation runs the given data through the model and does not perform any further actions. The outputs of this function will be collected for every batch and passed to the collate_predictions function.

Note

When using the predict function, this method is called within a torch.no_grad() block.

Parameters

data (object) – The batch as sampled from the data loader.

Returns

The return value passed to collate_predictions.

Return type

object

collate_losses(losses)

Combines the losses obtained from the train_batch() method. The default implementation assumes that simple floats are returned. Overwriting this method might be useful if you have different types of losses, such as generator and critic loss when training Wasserstein GANs.

Note

Overwrite this function if it is important to weigh losses obtained from batches in some manner (e.g. if data batches have different size).

Parameters

losses (list of object) – The losses returned from train_batch.

Returns

The loss names mapped to their values (usually float values).

Return type

dict of str -> object

collate_evals(evals)

Combines the evaluation objects returned from the eval_batch() method. The default implementation works whenever the returned objects are tensors, tuples of tensors or dictionaries of tensors. The contents will then simply be concatenated.

Parameters

evals (list of objects) – The evaluation objects.

Returns

An object that is passed to all metrics. If a tuple or a dict, they are passed as variadic arguments or keyword arguments, respectively.

Return type

object

collate_predictions(predictions)

Combines the predictions obtained for multiple batches as per the predict_batch() method. The default implementation works whenever the returned objects are tensors, tuples of tensors or dictionaries of tensors. The contents will then simply be concatenated.

Parameters

predictions (list of objects) – The predictions obtained.

Returns

An object that is obtained from the predictions.

Return type

object

to_device(device, item)

This method moves the given item to the provided device. The default implementation allows for Tensors as well as dicts, lists and tuples of tensors (they may be nested). If you require passing custom datatypes to a device as you feed them to your model or your model outputs them, you may overwrite this method.

Parameters

• device (torch.Device or str) – The device onto which the given item should be moved. May either be an actual device object (such as a GPU) or a string identifying the device (e.g. ‘cpu’ or ‘cuda:0’).

• item (object) – The item to move.

Returns

The item moved to the given device.

Return type

object

Maximum Likelihood Estimation

class pyblaze.nn.MLEEngine(model, expects_data_target=True, uses_data_target=True, uses_data_input=False)

This engine can be used for all kinds of learning tasks where some kind of likelihood is maximized. This includes both of supervised and unsupervised learning. In general, this engine should be used whenever a model simply optimizes a loss function.

The engine requires data to be available in the following format:

Parameters

• x (object) – The input to the model. The type depends on the model. Usually, however, it is a tensor of shape [N, …] for batch size N. This value is always required.

• y (object) – The labels corresponding to the input. The type depends on the model. Usually, it is set to the class indices (for classification) or the target values (for regression) and therefore given as tensor of shape [N]. In case this value is given as tuple, all components are given to the loss function. This is e.g. useful when using a triplet loss. The targets might or might not be given for unsupervised learning tasks. In the latter case, expects_data_target should be set to False in the engine’s initializer. Targets must not be given if predict() is called.

For the train() method allows for the following keyword arguments:

Parameters

• optimizer (torch.optim.Optimizer) – The optimizer to use for the model.

• loss (torch.nn.Module) – The loss function to use. The number of inputs it receives depends on the model being trained and the kind of learning problem. In the most general form, the parameters are given as (<all model outputs…>, <all targets…>, <all inputs…>) where the model output is expected to be a tuple. The passing of targets and inputs to the loss function may be toggled via the uses_data_target and uses_data_input parameters of the engine’s initializer. In the most common (and default) case, inputs are simply given as (torch.Tensor [N], torch.Tensor [N]) for batch size N. The output of this function must always be a torch.Tensor [1].

• gradient_accumulation_steps (int, default: 1) – The number of batches which should be used for a single update step. Gradient accumulation can be useful if the GPU can only fit a small batch size but model convergence is hindered by that. The number of gradient accumulation steps may not be changed within an epoch.

• kwargs (keyword arguments) – Additional keyword arguments passed to the forward() method of the model. These keyword arguments may also be passed for the eval_batch() method call.

Note

The same parameters that are given to the loss are also given to any metrics used with this engine.

__init__(model, expects_data_target=True, uses_data_target=True, uses_data_input=False)

Initializes a new MLE engine.

Parameters

• model (torch.nn.Module) – The model to train or evaluate.

• expects_data_target (bool, default: True) – Whether the data supplied to this engine is a tuple consisting of data and target or only contains the data. Setting this to False is never required for supervised learning.

• uses_data_target (bool, default: True) – Whether the loss function requires the target. Setting this to False although a data target is expected might make sense when the dataset comes with targets and it is easier to not preprocess it. This value is automatically set to False when no target is expected.

• uses_data_input (bool, default: True) – Whether the loss function requires the input to the model. This is e.g. useful for unsupervised learning tasks.

predict_batch(data, **kwargs)

Processes a single batch of inputs to make model predictions. The default implementation runs the given data through the model and does not perform any further actions. The outputs of this function will be collected for every batch and passed to the collate_predictions function.

Note

When using the predict function, this method is called within a torch.no_grad() block.

Parameters

data (object) – The batch as sampled from the data loader.

Returns

The return value passed to collate_predictions.

Return type

object

Autoencoders

class pyblaze.nn.AutoencoderEngine(model, expects_data_target=False)

Utility class for easily initializing the pyblaze.nn.MLEEngine for a usage with an autoencoder. This includes variational autoencoders.

When making use of the predict() method, the model passed to this engine should have a submodule named decoder. The predict() method takes the following additional parameters:

Parameters

reconstruct (bool, default: False) – If this flag is set, data passed to the predict() method is assumed to be the input to the encoder. Otherwise, the input is assumed to be from the latent space, i.e. to be the direct input to the decoder. In both cases, the output of the decoder is returned to the caller.

__init__(model, expects_data_target=False)

Initializes a new engine for autoencoders.

Parameters

• model (torch.nn.Module) – The model to train or evaluate.

• expects_data_target (bool, default: False) – Whether the data supplied to this engine is a tuple consisting of data and target or only contains the data.

predict_batch(data, reconstruct=False)

Processes a single batch of inputs to make model predictions. The default implementation runs the given data through the model and does not perform any further actions. The outputs of this function will be collected for every batch and passed to the collate_predictions function.

Note

When using the predict function, this method is called within a torch.no_grad() block.

Parameters

data (object) – The batch as sampled from the data loader.

Returns

The return value passed to collate_predictions.

Return type

object

Wasserstein GANs

class pyblaze.nn.WGANEngine(model, expects_data_target=False)

Engine to be used for training a Wasserstein GAN. It enables training the critic for multiple iterations as well as skipping the training of the generator for specific iterations.

The GAN must be supplied as a single model with generator and critic attributes. The output of the generator must be feedable directly to the forward() method of the critic. The supplied combined model does therefore not necessarily need to implement a forward() method.

The engine requires data to be available in the following format:

Parameters

• noise (list of object) – List of n + 1 noise inputs for the generator (batch size N, noise dimension D). When performing predictions, there is no need to pass a list.

• real (list of object) – List of n real inputs for the critic. Shape depends on the task at hand. When performing predictions, this value must not be given.

The train() method allows for the following keyword arguments:

Parameters

• generator_optimizer (torch.optim.Optimizer) – The optimizer to use for the generator model.

• critic_optimizer (torch.optim.Optimizer) – The optimizer to use for the critic model.

• generator_loss (torch.nn.Module) – The loss to use for the generator. Receives a tuple of the critic’s outputs for fake and real data (fake_out, real_out) (both values of type torch.Tensor [N] for batch size N) and ought to return a single torch.Tensor [1], the loss. Consider using pyblaze.nn.WassersteinLossGenerator.

• critic_loss (torch.nn.Module) – The loss to use for the critic. Receives a 4-tuple of the critic’s outputs for fake and real data as well as the the inputs to the critic (fake_out, real_out, fake_in, real_in) (the first two values are of type torch.Tensor [N], the second two values of type object depending on the data). It ought to return a tuple of (torch.Tensor [1], torch.Tensor [1]) where only the first value requires a gradient, the actual loss. The second value ought to provide an estimation of the earth mover’s distance (i.e. the negative loss when not penalizing gradients). Consider using pyblaze.nn.WassersteinLossCritic.

• critic_iterations (int, default: 3) – The number of iterations to train the critic for every iteration the generator is trained. A higher value ensures that the critic is trained to optimality and provides better gradients to the generator. However, it slows down training.

• skip_generator (bool, default: False) – Whether to skip optimizing the generator. Usually, you would want to alter this value during training via some callback. This way, you can ensure that the critic is trained to optimality at specific milestones.

• clip_weights (tuple of (float, float), default: None) – When this value is given, all weights of the generator are clamped into the given range after every step of the optimizer. While this ensures Lipschitz-continuity, using different means of regularization should be preferred.

• kwargs (keyword arguments) – Additional keyword arguments. If they start with ‘generator_’, they will be passed only to the generator, the same applies to keywords starting with ‘critic_’. When passed to the model, the prefixes are dropped.

Note

Calling evaluate() on this engine is not possible as there is no principled way to evaluate generator performance for an arbitrary generator. Subclass this engine to add the possibility for evaluation.

__init__(model, expects_data_target=False)

Initializes a new engine for training a WGAN.

Parameters

• model (torch.nn.Module) – The WGAN to train.

• expects_data_target (bool, default: False) – When this value is set to True, the real data instances passed to this engine are expected to yield class labels. These are simply discarded.

predict_batch(noise)

Processes a single batch of inputs to make model predictions. The default implementation runs the given data through the model and does not perform any further actions. The outputs of this function will be collected for every batch and passed to the collate_predictions function.

Note

When using the predict function, this method is called within a torch.no_grad() block.

Parameters

data (object) – The batch as sampled from the data loader.

Returns

The return value passed to collate_predictions.

Return type

object