import enum
import functools
import logging
import math
import time
from abc import ABC, abstractmethod
from collections import OrderedDict
from enum import Enum
from typing import List, Union, Sequence, Callable, TYPE_CHECKING, Tuple, Optional, Dict, Any
import matplotlib.figure
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from matplotlib import pyplot as plt
from torch import cuda as torchcuda
from .torch_data import TensorScaler, DataUtil, TorchDataSet, TorchDataSetProviderFromDataUtil, TorchDataSetProvider, \
TensorScalerIdentity, TensorTransformer
from .torch_enums import ClassificationOutputMode
from ..util.string import ToStringMixin
if TYPE_CHECKING:
from .torch_base import TorchModel
log = logging.getLogger(__name__)
[docs]class Optimiser(enum.Enum):
SGD = ("sgd", optim.SGD)
ASGD = ("asgd", optim.ASGD)
ADAGRAD = ("adagrad", optim.Adagrad)
ADADELTA = ("adadelta", optim.Adadelta)
ADAM = ("adam", optim.Adam)
ADAMW = ("adamw", optim.AdamW)
ADAMAX = ("adamax", optim.Adamax)
RMSPROP = ("rmsprop", optim.RMSprop)
RPROP = ("rprop", optim.Rprop)
LBFGS = ("lbfgs", optim.LBFGS)
[docs] @classmethod
def from_name(cls, name: str) -> "Optimiser":
lname = name.lower()
for o in cls:
if o.value[0] == lname:
return o
raise ValueError(f"Unknown optimiser name '{name}'; known names: {[o.value[0] for o in cls]}")
[docs] @classmethod
def from_name_or_instance(cls, name_or_instance: Union[str, "Optimiser"]) -> "Optimiser":
if type(name_or_instance) == str:
return cls.from_name(name_or_instance)
else:
return name_or_instance
class _Optimiser(object):
"""
Wrapper for classes inherited from torch.optim.Optimizer
"""
def __init__(self, params, method: Union[str, Optimiser], lr, max_grad_norm, use_shrinkage=True, **optimiser_args):
"""
:param params: an iterable of torch.Tensor s or dict s. Specifies what Tensors should be optimized.
:param method: the optimiser to use
:param lr: learnig rate
:param max_grad_norm: gradient norm value beyond which to apply gradient shrinkage
:param optimiser_args: keyword arguments to be used in actual torch optimiser
"""
self.method = Optimiser.from_name_or_instance(method)
self.params = list(params) # careful: params may be a generator
self.last_ppl = None
self.lr = lr
self.max_grad_norm = max_grad_norm
self.start_decay = False
self.optimiserArgs = optimiser_args
self.use_shrinkage = use_shrinkage
# instantiate optimiser
optimiser_args = dict(self.optimiserArgs)
optimiser_args.update({'lr': self.lr})
if self.method == Optimiser.LBFGS:
self.use_shrinkage = False
self.optimizer = optim.LBFGS(self.params, **optimiser_args)
else:
cons = self.method.value[1]
self.optimizer = cons(self.params, **optimiser_args)
def step(self, loss_backward: Callable):
"""
:param loss_backward: callable, performs backward step and returns loss
:return: loss value
"""
if self.use_shrinkage:
def closure_with_shrinkage():
loss_value = loss_backward()
torch.nn.utils.clip_grad_norm_(self.params, self.max_grad_norm)
return loss_value
closure = closure_with_shrinkage
else:
closure = loss_backward
loss = self.optimizer.step(closure)
return loss
[docs]class NNLossEvaluator(ABC):
"""
Base class defining the interface for training and validation loss evaluation.
"""
[docs] class Evaluation(ABC):
[docs] @abstractmethod
def start_epoch(self) -> None:
"""
Starts a new epoch, resetting any aggregated values required to ultimately return the
epoch's overall training loss (via getEpochTrainLoss) and validation metrics (via getValidationMetrics)
"""
pass
[docs] @abstractmethod
def compute_train_batch_loss(self, model_output, ground_truth, x, y) -> torch.Tensor:
"""
Computes the loss for the given model outputs and ground truth values for a batch
and aggregates the computed loss values such that :meth:``getEpochTrainLoss`` can return an appropriate
result for the entire epoch.
The original batch tensors X and Y are provided as meta-information only.
:param model_output: the model output
:param ground_truth: the ground truth values
:param x: the original batch input tensor
:param y: the original batch output (ground truth) tensor
:return: the loss (scalar tensor)
"""
pass
[docs] @abstractmethod
def get_epoch_train_loss(self) -> float:
"""
:return: the epoch's overall training loss (as obtained by collecting data from individual training
batch data passed to computeTrainBatchLoss)
"""
pass
[docs] @abstractmethod
def process_validation_batch(self, model_output, ground_truth, x, y) -> None:
"""
Processes the given model outputs and ground truth values in order to compute sufficient statistics for
velidation metrics, which at the end of the epoch, shall be retrievable via method getValidationMetrics
:param model_output: the model output
:param ground_truth: the ground truth values
:param x: the original batch input tensor
:param y: the original batch output (ground truth) tensor
:return: the loss (scalar tensor)
"""
pass
[docs] @abstractmethod
def get_validation_metrics(self) -> Dict[str, float]:
pass
[docs] @abstractmethod
def start_evaluation(self, cuda: bool) -> Evaluation:
"""
Begins the evaluation of a model, returning a (stateful) object which is to perform the necessary computations.
:param cuda: whether CUDA is being applied (all tensors/models on the GPU)
:return: the evaluation object
"""
pass
[docs] @abstractmethod
def get_validation_metric_name(self) -> str:
"""
:return: the name of the validation metric which is to be used to determine the best model (key for the ordered
dictionary returned by method Evaluation.getValidationMetrics)
"""
pass
[docs]class NNLossEvaluatorFixedDim(NNLossEvaluator, ABC):
"""
Base class defining the interface for training and validation loss evaluation, which uses fixed-dimension
outputs and aggregates individual training batch losses that are summed losses per batch
(averaging appropriately internally).
"""
[docs] class Evaluation(NNLossEvaluator.Evaluation):
def __init__(self, criterion, validation_loss_evaluator: "NNLossEvaluatorFixedDim.ValidationLossEvaluator",
output_dim_weights: torch.Tensor = None):
self.output_dim_weights = output_dim_weights
self.output_dim_weight_sum = torch.sum(output_dim_weights) if output_dim_weights is not None else None
self.validation_loss_evaluator = validation_loss_evaluator
self.criterion = criterion
self.total_loss = None
self.num_samples = None
self.num_outputs_per_data_point: Optional[int] = None
self.validation_ground_truth_shape = None
[docs] def start_epoch(self):
self.total_loss = 0
self.num_samples = 0
self.validation_ground_truth_shape = None
[docs] def compute_train_batch_loss(self, model_output, ground_truth, x, y) -> torch.Tensor:
# size of modelOutput and groundTruth: (batchSize, outputDim=numOutputsPerDataPoint)
if self.num_outputs_per_data_point is None:
output_shape = y.shape[1:]
self.num_outputs_per_data_point = functools.reduce(lambda x, y: x * y, output_shape, 1)
assert self.output_dim_weights is None or len(self.output_dim_weights) == self.num_outputs_per_data_point
num_data_points_in_batch = y.shape[0]
if self.output_dim_weights is None:
# treat all dimensions as equal, applying criterion to entire tensors
loss = self.criterion(model_output, ground_truth)
self.num_samples += num_data_points_in_batch * self.num_outputs_per_data_point
self.total_loss += loss.item()
return loss
else:
# compute loss per dimension and return weighted loss
loss_per_dim = torch.zeros(self.num_outputs_per_data_point, device=model_output.device, dtype=torch.float)
for o in range(self.num_outputs_per_data_point):
loss_per_dim[o] = self.criterion(model_output[:, o], ground_truth[:, o])
weighted_loss = (loss_per_dim * self.output_dim_weights).sum() / self.output_dim_weight_sum
self.num_samples += num_data_points_in_batch
self.total_loss += weighted_loss.item()
return weighted_loss
[docs] def get_epoch_train_loss(self) -> float:
return self.total_loss / self.num_samples
[docs] def process_validation_batch(self, model_output, ground_truth, x, y):
if self.validation_ground_truth_shape is None:
self.validation_ground_truth_shape = y.shape[1:] # the shape of the output of a single model application
self.validation_loss_evaluator.start_validation_collection(self.validation_ground_truth_shape)
self.validation_loss_evaluator.process_validation_result_batch(model_output, ground_truth)
[docs] def get_validation_metrics(self) -> Dict[str, float]:
return self.validation_loss_evaluator.end_validation_collection()
[docs] def start_evaluation(self, cuda: bool) -> Evaluation:
criterion = self.get_training_criterion()
output_dim_weights_array = self.get_output_dim_weights()
output_dim_weights_tensor = torch.from_numpy(output_dim_weights_array).float() if output_dim_weights_array is not None else None
if cuda:
criterion = criterion.cuda()
if output_dim_weights_tensor is not None:
output_dim_weights_tensor = output_dim_weights_tensor.cuda()
return self.Evaluation(criterion, self.create_validation_loss_evaluator(cuda), output_dim_weights=output_dim_weights_tensor)
[docs] @abstractmethod
def get_training_criterion(self) -> nn.Module:
"""
Gets the optimisation criterion (loss function) for training.
Standard implementations are available in torch.nn (torch.nn.MSELoss, torch.nn.CrossEntropyLoss, etc.).
"""
pass
[docs] @abstractmethod
def get_output_dim_weights(self) -> Optional[np.ndarray]:
pass
[docs] @abstractmethod
def create_validation_loss_evaluator(self, cuda: bool) -> "ValidationLossEvaluator":
"""
:param cuda: whether to use CUDA-based tensors
:return: the evaluator instance which is to be used to evaluate the model on validation data
"""
pass
[docs] def get_validation_metric_name(self) -> str:
"""
Gets the name of the metric (key of dictionary as returned by the validation loss evaluator's
endValidationCollection method), which is defining for the quality of the model and thus determines which
epoch's model is considered the best.
:return: the name of the metric
"""
pass
[docs] class ValidationLossEvaluator(ABC):
[docs] @abstractmethod
def start_validation_collection(self, ground_truth_shape):
"""
Initiates validation data collection for a new epoch, appropriately resetting this object's internal state.
:param ground_truth_shape: the tensor shape of a single ground truth data point (not including the batch
entry dimension)
"""
pass
[docs] @abstractmethod
def process_validation_result_batch(self, output, ground_truth):
"""
Collects, for validation, the given output and ground truth data (tensors holding data on one batch,
where the first dimension is the batch entry)
:param output: the model's output
:param ground_truth: the corresponding ground truth
"""
pass
[docs] @abstractmethod
def end_validation_collection(self) -> OrderedDict:
"""
Computes validation metrics based on the data previously processed.
:return: an ordered dictionary with validation metrics
"""
pass
[docs]class NNLossEvaluatorRegression(NNLossEvaluatorFixedDim, ToStringMixin):
"""A loss evaluator for (multi-variate) regression."""
[docs] class LossFunction(Enum):
L1LOSS = "L1Loss"
L2LOSS = "L2Loss"
MSELOSS = "MSELoss"
SMOOTHL1LOSS = "SmoothL1Loss"
def __init__(self, loss_fn: LossFunction = LossFunction.L2LOSS, validation_tensor_transformer: Optional[TensorTransformer] = None,
output_dim_weights: Sequence[float] = None, apply_output_dim_weights_in_validation=True,
validation_metric_name: Optional[str] = None):
"""
:param loss_fn: the loss function to use
:param validation_tensor_transformer: a transformer which is to be applied to validation tensors (both model outputs and ground
truth) prior to computing the validation metrics
:param output_dim_weights: vector of weights to apply to then mean loss per output dimension, i.e. for the case where for each data
point, the model produces n output dimensions, the mean loss for the i-th dimension is to be computed separately and be scaled
with the weight, and the overall loss returned is the weighted average. The weights need not sum to 1 (normalisation is
applied).
:param apply_output_dim_weights_in_validation: whether output dimension weights are also to be applied to to the metrics computed
for validation. Note that this may not be possible if a validationTensorTransformer which changes the output dimensions is
used.
:param validation_metric_name: the metric to use for model selection during validation; if None, use default depending on lossFn
"""
self.validation_tensor_transformer = validation_tensor_transformer
self.output_dim_weights = np.array(output_dim_weights) if output_dim_weights is not None else None
self.apply_output_dim_weights_in_validation = apply_output_dim_weights_in_validation
self.validation_metric_name = validation_metric_name
if loss_fn is None:
loss_fn = self.LossFunction.L2LOSS
try:
self.loss_fn = self.LossFunction(loss_fn)
except ValueError:
raise Exception(f"The loss function '{loss_fn}' is not supported. "
f"Available options are: {[e.value for e in self.LossFunction]}")
[docs] def create_validation_loss_evaluator(self, cuda):
return self.ValidationLossEvaluator(cuda, self.validation_tensor_transformer, self.output_dim_weights,
self.apply_output_dim_weights_in_validation)
[docs] def get_training_criterion(self):
if self.loss_fn is self.LossFunction.L1LOSS:
criterion = nn.L1Loss(reduction='sum')
elif self.loss_fn is self.LossFunction.L2LOSS or self.loss_fn == self.LossFunction.MSELOSS:
criterion = nn.MSELoss(reduction='sum')
elif self.loss_fn is self.LossFunction.SMOOTHL1LOSS:
criterion = nn.SmoothL1Loss(reduction='sum')
else:
raise AssertionError(f"Loss function {self.loss_fn} defined but instantiation not implemented.")
return criterion
[docs] def get_output_dim_weights(self) -> Optional[np.ndarray]:
return self.output_dim_weights
[docs] class ValidationLossEvaluator(NNLossEvaluatorFixedDim.ValidationLossEvaluator):
def __init__(self, cuda: bool, validation_tensor_transformer: Optional[TensorTransformer], output_dim_weights: np.ndarray,
apply_output_dim_weights: bool):
self.validationTensorTransformer = validation_tensor_transformer
self.outputDimWeights = output_dim_weights
self.applyOutputDimWeights = apply_output_dim_weights
self.total_loss_l1 = None
self.total_loss_l2 = None
self.output_dims = None
self.allTrueOutputs = None
self.evaluate_l1 = nn.L1Loss(reduction='sum')
self.evaluate_l2 = nn.MSELoss(reduction='sum')
if cuda:
self.evaluate_l1 = self.evaluate_l1.cuda()
self.evaluate_l2 = self.evaluate_l2.cuda()
self.begin_new_validation_collection: Optional[bool] = None
[docs] def start_validation_collection(self, ground_truth_shape):
if len(ground_truth_shape) != 1:
raise ValueError("Outputs that are not vectors are currently unsupported")
self.begin_new_validation_collection = True
[docs] def process_validation_result_batch(self, output, ground_truth):
# apply tensor transformer (if any)
if self.validationTensorTransformer is not None:
output = self.validationTensorTransformer.transform(output)
ground_truth = self.validationTensorTransformer.transform(ground_truth)
# check if new collection
if self.begin_new_validation_collection:
self.output_dims = ground_truth.shape[-1]
self.total_loss_l1 = np.zeros(self.output_dims)
self.total_loss_l2 = np.zeros(self.output_dims)
self.allTrueOutputs = None
self.begin_new_validation_collection = False
assert len(output.shape) == 2 and len(ground_truth.shape) == 2
# obtain series of outputs per output dimension: (batch_size, output_size) -> (output_size, batch_size)
predicted_output = output.permute(1, 0)
true_output = ground_truth.permute(1, 0)
if self.allTrueOutputs is None:
self.allTrueOutputs = true_output
else:
self.allTrueOutputs = torch.cat((self.allTrueOutputs, true_output), dim=1)
for i in range(self.output_dims):
self.total_loss_l1[i] += self.evaluate_l1(predicted_output[i], true_output[i]).item()
self.total_loss_l2[i] += self.evaluate_l2(predicted_output[i], true_output[i]).item()
[docs] def end_validation_collection(self):
output_dims = self.output_dims
rae = np.zeros(output_dims)
rrse = np.zeros(output_dims)
mae = np.zeros(output_dims)
mse = np.zeros(output_dims)
for i in range(output_dims):
mean = torch.mean(self.allTrueOutputs[i])
ref_model_errors = self.allTrueOutputs[i] - mean
ref_model_sum_abs_errors = torch.sum(torch.abs(ref_model_errors)).item()
ref_model_sum_squared_errors = torch.sum(ref_model_errors * ref_model_errors).item()
num_samples = ref_model_errors.size(0)
mae[i] = self.total_loss_l1[i] / num_samples
mse[i] = self.total_loss_l2[i] / num_samples
rae[i] = self.total_loss_l1[i] / ref_model_sum_abs_errors if ref_model_sum_abs_errors != 0 else np.inf
rrse[i] = np.sqrt(mse[i]) / np.sqrt(
ref_model_sum_squared_errors / num_samples) if ref_model_sum_squared_errors != 0 else np.inf
def mean(x):
if self.applyOutputDimWeights:
return np.average(x, weights=self.outputDimWeights)
else:
return np.mean(x)
metrics = OrderedDict([("RRSE", mean(rrse)), ("RAE", mean(rae)), ("MSE", mean(mse)), ("MAE", mean(mae))])
return metrics
[docs] def get_validation_metric_name(self):
if self.validation_metric_name is not None:
return self.validation_metric_name
else:
if self.loss_fn is self.LossFunction.L1LOSS or self.loss_fn is self.LossFunction.SMOOTHL1LOSS:
return "MAE"
elif self.loss_fn is self.LossFunction.L2LOSS or self.loss_fn is self.LossFunction.MSELOSS:
return "MSE"
else:
raise AssertionError(f"No validation metric defined as selection criterion for loss function {self.loss_fn}")
[docs]class NNLossEvaluatorClassification(NNLossEvaluatorFixedDim):
"""A loss evaluator for classification"""
[docs] class LossFunction(Enum):
CROSSENTROPY = "CrossEntropy"
NLL = "NegativeLogLikelihood"
[docs] def create_criterion(self) -> Callable:
if self is self.CROSSENTROPY:
return nn.CrossEntropyLoss(reduction='sum')
elif self is self.NLL:
return nn.NLLLoss(reduction="sum")
[docs] def get_validation_metric_key(self) -> str:
if self is self.CROSSENTROPY:
return "CE"
elif self is self.NLL:
return "NLL"
[docs] @classmethod
def default_for_output_mode(cls, output_mode: ClassificationOutputMode):
if output_mode == ClassificationOutputMode.PROBABILITIES:
raise ValueError(f"No loss function available for {output_mode}; Either apply log at the end and use "
f"{ClassificationOutputMode.LOG_PROBABILITIES} or use a different final activation (e.g. log_softmax) "
f"to avoid this type of output.")
elif output_mode == ClassificationOutputMode.LOG_PROBABILITIES:
return cls.NLL
elif output_mode == ClassificationOutputMode.UNNORMALISED_LOG_PROBABILITIES:
return cls.CROSSENTROPY
else:
raise ValueError(f"No default specified for {output_mode}")
def __init__(self, loss_fn: LossFunction):
self.lossFn: "NNLossEvaluatorClassification.LossFunction" = self.LossFunction(loss_fn)
def __str__(self):
return f"{self.__class__.__name__}[{self.lossFn}]"
[docs] def create_validation_loss_evaluator(self, cuda):
return self.ValidationLossEvaluator(cuda, self.lossFn)
[docs] def get_training_criterion(self):
return self.lossFn.create_criterion()
[docs] def get_output_dim_weights(self) -> Optional[np.ndarray]:
return None
[docs] class ValidationLossEvaluator(NNLossEvaluatorFixedDim.ValidationLossEvaluator):
def __init__(self, cuda: bool, loss_fn: "NNLossEvaluatorClassification.LossFunction"):
self.loss_fn = loss_fn
self.total_loss = None
self.num_validation_samples = None
self.criterion = self.loss_fn.create_criterion()
if cuda:
self.criterion = self.criterion.cuda()
[docs] def start_validation_collection(self, ground_truth_shape):
self.total_loss = 0
self.num_validation_samples = 0
[docs] def process_validation_result_batch(self, output, ground_truth):
self.total_loss += self.criterion(output, ground_truth).item()
self.num_validation_samples += output.shape[0]
[docs] def end_validation_collection(self):
mean_loss = self.total_loss / self.num_validation_samples
if isinstance(self.criterion, nn.CrossEntropyLoss):
metrics = OrderedDict([("CE", mean_loss), ("GeoMeanProbTrueClass", math.exp(-mean_loss))])
elif isinstance(self.criterion, nn.NLLLoss):
metrics = {"NLL": mean_loss}
else:
raise ValueError()
return metrics
[docs] def get_validation_metric_name(self):
return self.lossFn.get_validation_metric_key()
[docs]class NNOptimiserParams(ToStringMixin):
REMOVED_PARAMS = {"cuda"}
RENAMED_PARAMS = {
"optimiserClip": "optimiser_clip",
"lossEvaluator": "loss_evaluator",
"optimiserLR": "optimiser_lr",
"earlyStoppingEpochs": "early_stopping_epochs",
"batchSize": "batch_size",
"trainFraction": "train_fraction",
"scaledOutputs": "scaled_outputs",
"useShrinkage": "use_shrinkage",
"shrinkageClip": "shrinkage_clip",
}
def __init__(self,
loss_evaluator: NNLossEvaluator = None,
gpu: Optional[int] = None,
optimiser: Union[str, Optimiser] = "adam",
optimiser_lr=0.001,
early_stopping_epochs=None,
batch_size=None,
epochs=1000,
train_fraction=0.75,
scaled_outputs=False,
use_shrinkage=True,
shrinkage_clip=10.,
shuffle=True,
optimiser_args: Optional[Dict[str, Any]] = None):
"""
:param loss_evaluator: the loss evaluator to use
:param gpu: the index of the GPU to be used (if CUDA is enabled for the model to be trained); if None, default to first GPU
:param optimiser: the optimiser to use
:param optimiser_lr: the optimiser's learning rate
:param early_stopping_epochs: the number of epochs without validation score improvement after which to abort training and
use the best epoch's model (early stopping); if None, never abort training before all epochs are completed
:param batch_size: the batch size to use; for algorithms L-BFGS (optimiser='lbfgs'), which do not use batches, leave this at None.
If the algorithm uses batches and None is specified, batch size 64 will be used by default.
:param train_fraction: the fraction of the data used for training (with the remainder being used for validation).
If no validation is to be performed, pass 1.0.
:param scaled_outputs: whether to scale all outputs, resulting in computations of the loss function based on scaled values rather
than normalised values.
Enabling scaling may not be appropriate in cases where there are multiple outputs on different scales/with completely different
units.
:param use_shrinkage: whether to apply shrinkage to gradients whose norm exceeds ``shrinkageClip``, scaling the gradient down to
``shrinkageClip``
:param shrinkage_clip: the maximum gradient norm beyond which to apply shrinkage (if ``useShrinkage`` is True)
:param shuffle: whether to shuffle the training data
:param optimiser_args: keyword arguments to be passed on to the actual torch optimiser
"""
if Optimiser.from_name_or_instance(optimiser) == Optimiser.LBFGS:
large_batch_size = 1e12
if batch_size is not None:
log.warning(f"LBFGS does not make use of batches, therefore using large batch size {large_batch_size} "
f"to achieve use of a single batch")
batch_size = large_batch_size
else:
if batch_size is None:
log.debug("No batch size was specified, using batch size 64 by default")
batch_size = 64
self.epochs = epochs
self.batch_size = batch_size
self.optimiser_lr = optimiser_lr
self.shrinkage_clip = shrinkage_clip
self.optimiser = optimiser
self.gpu = gpu
self.train_fraction = train_fraction
self.scaled_outputs = scaled_outputs
self.loss_evaluator = loss_evaluator
self.optimiser_args = optimiser_args if optimiser_args is not None else {}
self.use_shrinkage = use_shrinkage
self.early_stopping_epochs = early_stopping_epochs
self.shuffle = shuffle
@classmethod
def _updated_params(cls, params: dict) -> dict:
return {cls.RENAMED_PARAMS.get(k, k): v for k, v in params.items() if k not in cls.REMOVED_PARAMS}
def __setstate__(self, state):
if "shuffle" not in state:
state["shuffle"] = True
self.__dict__ = self._updated_params(state)
[docs] @classmethod
def from_dict_or_instance(cls, nn_optimiser_params: Union[dict, "NNOptimiserParams"]) -> "NNOptimiserParams":
if isinstance(nn_optimiser_params, NNOptimiserParams):
return nn_optimiser_params
else:
return cls.from_dict(nn_optimiser_params)
[docs] @classmethod
def from_dict(cls, params: dict) -> "NNOptimiserParams":
return NNOptimiserParams(**cls._updated_params(params))
# TODO remove deprecated dict interface
[docs] @classmethod
def from_either_dict_or_instance(cls, nn_optimiser_dict_params: dict, nn_optimiser_params: Optional["NNOptimiserParams"]):
have_instance = nn_optimiser_params is not None
have_dict = len(nn_optimiser_dict_params)
if have_instance and have_dict:
raise ValueError("Received both a non-empty dictionary and an instance")
if have_instance:
return nn_optimiser_params
else:
return NNOptimiserParams.from_dict(nn_optimiser_dict_params)
[docs]class NNOptimiser:
log = log.getChild(__qualname__)
def __init__(self, params: NNOptimiserParams):
"""
:param params: parameters
"""
if params.loss_evaluator is None:
raise ValueError("Must provide a loss evaluator")
self.params = params
self.cuda = None
self.best_epoch = None
def __str__(self):
return f"{self.__class__.__name__}[params={self.params}]"
[docs] def fit(self,
model: "TorchModel",
data: Union[DataUtil, List[DataUtil], TorchDataSetProvider, List[TorchDataSetProvider],
TorchDataSet, List[TorchDataSet], Tuple[TorchDataSet, TorchDataSet], List[Tuple[TorchDataSet, TorchDataSet]]],
create_torch_module=True) -> "TrainingInfo":
"""
Fits the parameters of the given model to the given data, which can be a list of or single instance of one of the following:
* a `DataUtil` or `TorchDataSetProvider` (from which a training set and validation set will be obtained according to
the `trainFraction` parameter of this object)
* a `TorchDataSet` which shall be used as the training set (for the case where no validation set shall be used)
* a tuple with two `TorchDataSet` instances, where the first shall be used as the training set and the second as
the validation set
:param model: the model to be fitted
:param data: the data to use (see variants above)
:param create_torch_module: whether to newly create the torch module that is to be trained from the model's factory.
If False, (re-)train the existing module.
"""
self.cuda = model.cuda
self.log.info(f"Preparing parameter learning of {model} via {self} with cuda={self.cuda}")
use_validation = self.params.train_fraction != 1.0
def to_data_set_provider(d) -> TorchDataSetProvider:
if isinstance(d, TorchDataSetProvider):
return d
elif isinstance(d, DataUtil):
return TorchDataSetProviderFromDataUtil(d, self.cuda)
else:
raise ValueError(f"Cannot create a TorchDataSetProvider from {d}")
training_log_entries = []
def training_log(s):
self.log.info(s)
training_log_entries.append(s)
self._init_cuda()
# Set the random seed manually for reproducibility.
seed = 42
torch.manual_seed(seed)
if self.cuda:
torchcuda.manual_seed_all(seed)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
# obtain data, splitting it into training and validation set(s)
validation_sets = []
training_sets = []
output_scalers = []
if type(data) != list:
data = [data]
self.log.info("Obtaining input/output training instances")
for idx_data_item, data_item in enumerate(data):
if isinstance(data_item, TorchDataSet):
if use_validation:
raise ValueError("Passing a TorchDataSet instance is not admissible when validation is enabled (trainFraction != 1.0). "
"Pass a TorchDataSetProvider or another representation that supports validation instead.")
training_set = data_item
validation_set = None
output_scaler = TensorScalerIdentity()
elif type(data_item) == tuple:
training_set, validation_set = data_item
output_scaler = TensorScalerIdentity()
else:
data_set_provider = to_data_set_provider(data_item)
training_set, validation_set = data_set_provider.provide_split(self.params.train_fraction)
output_scaler = data_set_provider.get_output_tensor_scaler()
training_sets.append(training_set)
if validation_set is not None:
validation_sets.append(validation_set)
output_scalers.append(output_scaler)
training_log(f"Data set {idx_data_item+1}/{len(data)}: #train={training_set.size()}, "
f"#validation={validation_set.size() if validation_set is not None else 'None'}")
training_log("Number of validation sets: %d" % len(validation_sets))
torch_model = model.create_torch_module() if create_torch_module else model.get_torch_module()
if self.cuda:
torch_model.cuda()
model.set_torch_module(torch_model)
n_params = sum([p.nelement() for p in torch_model.parameters()])
self.log.info(f"Learning parameters of {model}")
training_log('Number of parameters: %d' % n_params)
training_log(f"Starting training process via {self}")
loss_evaluator = self.params.loss_evaluator
total_epochs = None
best_val = 1e9
best_epoch = 0
optim = _Optimiser(torch_model.parameters(), method=self.params.optimiser, lr=self.params.optimiser_lr,
max_grad_norm=self.params.shrinkage_clip, use_shrinkage=self.params.use_shrinkage, **self.params.optimiser_args)
best_model_bytes = model.get_module_bytes()
loss_evaluation = loss_evaluator.start_evaluation(self.cuda)
validation_metric_name = loss_evaluator.get_validation_metric_name()
training_loss_values = []
validation_metric_values = []
try:
self.log.info(f'Begin training with cuda={self.cuda}')
self.log.info('Press Ctrl+C to end training early')
for epoch in range(1, self.params.epochs + 1):
loss_evaluation.start_epoch()
epoch_start_time = time.time()
# perform training step, processing all the training data once
train_loss = self._train(training_sets, torch_model, optim, loss_evaluation, self.params.batch_size, output_scalers)
training_loss_values.append(train_loss)
# perform validation, computing the mean metrics across all validation sets (if more than one),
# and check for new best result according to validation results
is_new_best = False
if use_validation:
metrics_sum = None
metrics_keys = None
for i, (validation_set, output_scaler) in enumerate(zip(validation_sets, output_scalers)):
metrics = self._evaluate(validation_set, torch_model, loss_evaluation, output_scaler)
metrics_array = np.array(list(metrics.values()))
if i == 0:
metrics_sum = metrics_array
metrics_keys = metrics.keys()
else:
metrics_sum += metrics_array
metrics_sum /= len(validation_sets) # mean results
metrics = dict(zip(metrics_keys, metrics_sum))
current_val = metrics[loss_evaluator.get_validation_metric_name()]
validation_metric_values.append(current_val)
is_new_best = current_val < best_val
if is_new_best:
best_val = current_val
best_epoch = epoch
best_str = "best {:s} {:5.6f} from this epoch".format(validation_metric_name, best_val)
else:
best_str = "best {:s} {:5.6f} from epoch {:d}".format(validation_metric_name, best_val, best_epoch)
val_str = f' | validation {", ".join(["%s %5.4f" % e for e in metrics.items()])} | {best_str}'
else:
val_str = ""
training_log(
'Epoch {:3d}/{} completed in {:5.2f}s | train loss {:5.4f}{:s}'.format(
epoch, self.params.epochs, (time.time() - epoch_start_time), train_loss, val_str))
total_epochs = epoch
if use_validation:
if is_new_best:
best_model_bytes = model.get_module_bytes()
# check for early stopping
num_epochs_without_improvement = epoch - best_epoch
if self.params.early_stopping_epochs is not None and \
num_epochs_without_improvement >= self.params.early_stopping_epochs:
training_log(f"Stopping early: {num_epochs_without_improvement} epochs without validation metric improvement")
break
training_log("Training complete")
except KeyboardInterrupt:
training_log('Exiting from training early because of keyboard interrupt')
# reload best model according to validation results
if use_validation:
training_log(f'Best model is from epoch {best_epoch} with {validation_metric_name} {best_val} on validation set')
self.best_epoch = best_epoch
model.set_module_bytes(best_model_bytes)
return TrainingInfo(best_epoch=best_epoch if use_validation else None, log=training_log_entries, total_epochs=total_epochs,
training_loss_sequence=training_loss_values, validation_metric_sequence=validation_metric_values)
def _apply_model(self, model, input: Union[torch.Tensor, Sequence[torch.Tensor]], ground_truth, output_scaler: TensorScaler):
if isinstance(input, torch.Tensor):
output = model(input)
else:
output = model(*input)
if self.params.scaled_outputs:
output, ground_truth = self._scaled_values(output, ground_truth, output_scaler)
return output, ground_truth
@classmethod
def _scaled_values(cls, model_output, ground_truth, output_scaler):
scaled_output = output_scaler.denormalise(model_output)
scaled_truth = output_scaler.denormalise(ground_truth)
return scaled_output, scaled_truth
def _train(self, data_sets: Sequence[TorchDataSet], model: nn.Module, optim: _Optimiser,
loss_evaluation: NNLossEvaluator.Evaluation, batch_size: int, output_scalers: Sequence[TensorScaler]):
"""Performs one training epoch"""
model.train()
for data_set, output_scaler in zip(data_sets, output_scalers):
for X, Y in data_set.iter_batches(batch_size, shuffle=self.params.shuffle):
def closure():
model.zero_grad()
output, ground_truth = self._apply_model(model, X, Y, output_scaler)
loss = loss_evaluation.compute_train_batch_loss(output, ground_truth, X, Y)
loss.backward()
return loss
optim.step(closure)
return loss_evaluation.get_epoch_train_loss()
def _evaluate(self, data_set: TorchDataSet, model: nn.Module, loss_evaluation: NNLossEvaluator.Evaluation,
output_scaler: TensorScaler):
"""Evaluates the model on the given data set (a validation set)"""
model.eval()
for X, Y in data_set.iter_batches(self.params.batch_size, shuffle=False):
with torch.no_grad():
output, ground_truth = self._apply_model(model, X, Y, output_scaler)
loss_evaluation.process_validation_batch(output, ground_truth, X, Y)
return loss_evaluation.get_validation_metrics()
def _init_cuda(self):
"""Initialises CUDA (for learning) by setting the appropriate device if necessary"""
if self.cuda:
device_count = torchcuda.device_count()
if device_count == 0:
raise Exception("CUDA is enabled but no device found")
if self.params.gpu is None:
if device_count > 1:
log.warning("More than one GPU detected but no GPU index was specified, using GPU 0 by default.")
gpu_index = 0
else:
gpu_index = self.params.gpu
torchcuda.set_device(gpu_index)
elif torchcuda.is_available():
self.log.info("NOTE: You have a CUDA device; consider running with cuda=True")
[docs]class TrainingInfo:
def __init__(self, best_epoch: int = None, log: List[str] = None, training_loss_sequence: Sequence[float] = None,
validation_metric_sequence: Sequence[float] = None, total_epochs=None):
self.validation_metric_sequence = validation_metric_sequence
self.training_loss_sequence = training_loss_sequence
self.log = log
self.best_epoch = best_epoch
self.total_epochs = total_epochs
def __setstate__(self, state):
if "totalEpochs" not in state:
state["totalEpochs"] = None
self.__dict__ = state
[docs] def get_training_loss_series(self) -> pd.Series:
return self._create_series_with_one_based_index(self.training_loss_sequence, name="training loss")
[docs] def get_validation_metric_series(self) -> pd.Series:
return self._create_series_with_one_based_index(self.validation_metric_sequence, name="validation metric")
def _create_series_with_one_based_index(self, sequence: Sequence, name: str):
series = pd.Series(sequence, name=name)
series.index += 1
return series
[docs] def plot_all(self) -> matplotlib.figure.Figure:
"""
Plots both the sequence of training loss values and the sequence of validation metric values
"""
ts = self.get_training_loss_series()
vs = self.get_validation_metric_series()
fig, primary_ax = plt.subplots(1, 1)
secondary_ax = primary_ax.twinx()
training_line = primary_ax.plot(ts, color='blue')
validation_line = secondary_ax.plot(vs, color='orange')
best_epoc_line = primary_ax.axvline(self.best_epoch, color='black', linestyle='dashed')
primary_ax.set_xlabel("epoch")
primary_ax.set_ylabel(ts.name)
secondary_ax.set_ylabel(vs.name)
primary_ax.legend(training_line + validation_line + [best_epoc_line], [ts.name, vs.name, "best epoch"])
plt.tight_layout()
return fig
