class LogisticRegression:
# We can set/fix hyper
def __init__(self, random_state, **hyper_parameters):
self.model_ctor = linear_model.LogisticRegression
self.random_state = random_state
self.hp = HyperParameters(self.hyperparameter_space(),
**hyper_parameters)
self.model = None
# List of all the hyper-parameters
@staticmethod
def hyperparameter_space():
return {'C': 'uniform(0, 1)', 'l1_ratio': 'uniform(0, 1)'}
# List of hyper-parameters that needs to be set to finish initialization
def get_space(self):
return self.hp.missing_parameters()
# Initialize the model when all the hyper-parameters are there
def init(self, **hyper_parameters):
self.hp.add_parameters(**hyper_parameters)
self.model = self.model_ctor(penalty='elasticnet',
solver='saga',
random_state=self.random_state,
**self.hp.parameters(strict=True))
def predict(self, x):
try:
return self.model.predict(x)
except sklearn.exceptions.NotFittedError as e:
raise NotFittedError from e
def fit(self, x, y):
self.model = self.model.fit(x, y)
return self.model
class Initializer:
"""Lazy Initializer"""
def __init__(self, name, seed=0, **kwargs):
self.name = name
self.hyper_parameters = HyperParameters(space={})
self.seed = seed
self._initializer = None
self.initializer_ctor = registered_initialization.get(name)
if self.initializer_ctor is None:
raise RegisteredInitNotFound(name)
if hasattr(self.initializer_ctor, 'get_space'):
self.hyper_parameters.space = self.initializer_ctor.get_space()
self.hyper_parameters.add_parameters(**kwargs)
def get_space(self):
"""Return the dimension space of each parameters"""
return self.hyper_parameters.missing_parameters()
def get_current_space(self):
"""Get currently defined parameter space"""
return self.hyper_parameters.parameters(strict=False)
def init(self, override=False, **kwargs):
if self._initializer and not override:
warning(
'Initializer is already set, use override=True to force re initialization'
)
return self
self.hyper_parameters.add_parameters(**kwargs)
self._initializer = self.initializer_ctor(
**self.hyper_parameters.parameters(strict=True))
return self
@property
def initializer(self):
if not self._initializer:
self.init()
return self._initializer
def __call__(self, model):
with fork_rng(enabled=True):
init_seed(self.seed)
return self.initializer(model)
def test_hyperparameter_tracking():
space = {
'a': 'uniform(0, 1)',
'b': 'uniform(0, 1)'
}
# space with Fixed HP
hp = HyperParameters(space, b=0.124)
# Hp a is missing
with pytest.raises(MissingParameters):
hp.parameters(strict=True)
# return the space of missing params
assert hp.missing_parameters() == dict(a='uniform(0, 1)')
hp.add_parameters(a=0.123)
assert hp.missing_parameters() == {}
assert hp.parameters(strict=True) == dict(a=0.123, b=0.124)
class DecisionTree:
# We can set/fix hyper
def __init__(self, random_state, **hyper_parameters):
self.model_ctor = tree.DecisionTreeClassifier
self.random_state = random_state
self.hp = HyperParameters(self.hyperparameter_space(), **hyper_parameters)
self.model = None
# List of all the hyper-parameters
@staticmethod
def hyperparameter_space():
return {
'max_depth': 'uniform(0, 100, discrete=True)',
'min_samples_split': 'uniform(1, 10, discrete=True)',
'min_samples_leaf': 'uniform(1, 10, discrete=True)',
'min_weight_fraction_leaf': 'uniform(0, 1)'
}
# List of hyper-parameters that needs to be set to finish initialization
def get_space(self):
return self.hp.missing_parameters()
# Initialize the model when all the hyper-parameters are there
def init(self, **hyper_parameters):
self.hp.add_parameters(**hyper_parameters)
self.model = self.model_ctor(
criterion='gini',
splitter='best',
max_features=None,
random_state=self.random_state,
**self.hp.parameters(strict=True)
)
def predict(self, x):
try:
return self.model.predict(x)
except sklearn.exceptions.NotFittedError as e:
raise NotFittedError from e
def fit(self, x, y):
self.model = self.model.fit(x, y)
return self.model
def test_hyperparameter_nested_tracking():
space = {
'initializer': {
'a': 'uniform(0, 1)',
'b': 'uniform(0, 1)',
}
}
hp = HyperParameters(space, initializer=dict(b=0.124))
# Hp a is missing
with pytest.raises(MissingParameters):
hp.parameters(strict=True)
# return the space of missing params
assert hp.missing_parameters() == dict(initializer=dict(a='uniform(0, 1)'))
hp.add_parameters(initializer=dict(a=0.123))
assert hp.missing_parameters() == {}
assert hp.parameters(strict=True) == dict(initializer=dict(a=0.123, b=0.124))
class MLPRegressor:
# We can set/fix hyper
def __init__(self, random_state, **hyper_parameters):
self.model_ctor = sklearn.neural_network.MLPRegressor
self.random_state = random_state
self.hp = HyperParameters(self.hyperparameter_space(),
**hyper_parameters)
self.model = None
# List of all the hyper-parameters
@staticmethod
def hyperparameter_space():
return {
'hidden_layer_sizes': 'uniform(50, 70, discrete=True)',
'solver': 'uniform(0, 3, discrete=True)',
'alpha': 'uniform(0, 0.1)'
}
# List of hyper-parameters that needs to be set to finish initialization
def get_space(self):
return self.hp.missing_parameters()
# Initialize the model when all the hyper-parameters are there
def init(self, **hyper_parameters):
self.hp.add_parameters(**hyper_parameters)
self.model = self.model_ctor(random_state=self.random_state,
**self.hp.parameters(strict=True))
def predict(self, x):
#import pdb; pdb.set_trace()
#TODO: check if this fixes it
x = x[0]
try:
return self.model.predict(x)
except sklearn.exceptions.NotFittedError as e:
raise NotFittedError from e
def fit(self, x, y):
self.model = self.model.fit(x, y)
return self.model
class Optimizer(TorchOptimizer):
"""Lazy Optimizer that allows you to first fetch the supported parameters using ``get_space`` and then
initialize the underlying optimizer using ``init_optimizer``
Parameters
----------
name: str
Name of a registered optimizer
optimizer: Optimizer
Custom optimizer, mutually exclusive with :param name
half: bool
Enable fp16 Optimizer
loss_scale: float (LS)
fp16 optimizer option: loss scale to use
dynamic_loss_scale: bool
fp16 optimizer option: Enable dynamic loss scaling
scale_window: int (SW)
dynamic loss scaling option: Increase LS after SW successful iteration
scale_factor: float (SF)
dynamic loss scaling option: divide LS by SF after an overflow, or
multiply LS by SF after SW successful iteration
min_loss_scale: float
max_loss_scale: float
Examples
--------
Follows standard Pytorch Optimizer
>>> import torch
>>> from olympus.models import Model
>>> model = Model('resnet18',
... input_size=(1, 28, 28),
... output_size=10,)
>>>
>>> x = torch.randn((1, 1, 28, 28))
>>>
>>> optimizer = Optimizer('SGD', params=model.parameters(), weight_decay=1e-3, lr=0.001, momentum=0.8)
>>>
>>> optimizer.zero_grad()
>>> loss = model(x).sum()
>>> optimizer.backward(loss)
>>> optimizer.step()
Can be lazily initialized for hyper parameter search
>>> optimizer = Optimizer('SGD')
>>> optimizer.get_space()
{'lr': 'loguniform(1e-5, 1)', 'momentum': 'uniform(0, 1)', 'weight_decay': 'loguniform(1e-10, 1e-3)'}
>>> optimizer.init(model.parameters(), weight_decay=1e-3, lr=0.001, momentum=0.8)
>>>
>>> optimizer.zero_grad()
>>> loss = model(x).sum()
>>> optimizer.backward(loss)
>>> optimizer.step()
Switch to a mixed precision optimizer if needed
>>> optimizer = Optimizer('SGD', half=True)
Raises
------
RegisteredOptimizerNotFound
when using a name of an known optimizers
MissingArgument:
if name nor optimizer were not set
WrongParameter
if a wrong hyper parameter is passed in kwargs
"""
half = False
half_args = dict()
_optimizer = None
def __init__(self,
name=None,
*,
params=None,
optimizer=None,
half=False,
loss_scale=1,
dynamic_loss_scale=False,
scale_window=1000,
scale_factor=2,
min_loss_scale=None,
max_loss_scale=2.**24,
**kwargs):
self._optimizer = None
if params is not None:
params = list(params)
assert isinstance(params, (list, tuple))
self._model_parameters = params
self._half_parameters(half, loss_scale, dynamic_loss_scale,
scale_window, scale_factor, min_loss_scale,
max_loss_scale)
# Track defined hyper parameters
self.hyper_parameters = HyperParameters(space={})
if optimizer:
warning('Using custom optimizer')
if isinstance(optimizer, type):
self.optimizer_builder = optimizer
if hasattr(optimizer, 'get_space'):
self.hyper_parameters.space = optimizer.get_space()
else:
self._optimizer = self._wrap_optimizer(optimizer)
if hasattr(self._optimizer, 'get_space'):
self.hyper_parameters.space = self._optimizer.get_space()
elif name:
# load an olympus model
self.optimizer_builder = registered_optimizers.get(name.lower())
if not self.optimizer_builder:
raise RegisteredOptimizerNotFound(name)
if hasattr(self.optimizer_builder, 'get_space'):
self.hyper_parameters.space = self.optimizer_builder.get_space(
)
else:
raise MissingArgument('optimizer or name needs to be set')
# All additional args are hyper parameters
self.hyper_parameters.add_parameters(**kwargs)
def _half_parameters(self,
half=False,
loss_scale=1,
dynamic_loss_scale=False,
scale_window=1000,
scale_factor=2,
min_loss_scale=None,
max_loss_scale=2.**24):
"""Save the configuration of the fp16 optimizer"""
self.half = half
static_loss_scale = loss_scale
if dynamic_loss_scale:
static_loss_scale = 'dynamic'
self.half_args = dict(static_loss_scale=static_loss_scale,
dynamic_loss_scale=dynamic_loss_scale,
dynamic_loss_args=dict(
init_scale=loss_scale,
scale_factor=scale_factor,
scale_window=scale_window,
min_loss_scale=min_loss_scale,
max_loss_scale=max_loss_scale),
verbose=False)
def _wrap_optimizer(self, optimizer):
if self.half:
from olympus.utils.fp16 import FP16Optimizer
return FP16Optimizer(optimizer, **self.half_args)
return optimizer
def get_space(self) -> Dict[str, str]:
"""Return the dimension space of each parameters"""
if self._optimizer:
warning('Optimizer is already set')
return self.hyper_parameters.missing_parameters()
def get_current_space(self):
"""Get currently defined parameter space"""
return self.hyper_parameters.parameters(strict=False)
@property
def defaults(self):
"""Returns the default hyper parameter of the underlying optimizer"""
return self.optimizer_builder.defaults()
def init(self, params=None, override=False, **kwargs):
"""instantiate the underlying optimizer
Raises
------
MissingParameters
if an hyper parameter is missing
"""
if params is not None:
params = list(params)
assert isinstance(params, (list, tuple))
if self._optimizer and not override:
warning(
'Optimizer is already set, use override=True to force re initialization'
)
return self
# add missing hyper parameters
self.hyper_parameters.add_parameters(**kwargs)
if params is None:
params = self._model_parameters
if params is None:
raise MissingArgument('Missing Model parameters!')
self._optimizer = self._wrap_optimizer(
self.optimizer_builder(
params, **self.hyper_parameters.parameters(strict=True)))
@property
def optimizer(self):
if not self._optimizer:
self.init()
return self._optimizer
def backward(self, loss):
if self.half: # for loss scaling
self.optimizer.backward(loss)
else:
loss.backward()
def step(self, closure=None):
return self.optimizer.step(closure)
def zero_grad(self):
return self.optimizer.zero_grad()
@property
def param_groups(self):
return self.optimizer.param_groups
@property
def state(self):
return self.optimizer.state
def to(self, device):
if self._optimizer:
for state in self.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.to(device=device)
return self
def state_dict(self, destination=None, prefix='', keep_vars=False):
s = self.optimizer.state_dict()
return s
def load_state_dict(self, state_dict, strict=True, device=None):
self.optimizer.load_state_dict(state_dict)
if device:
self.to(device)
class LRSchedule:
"""Lazy LRSchedule that allows you to first fetch the supported parameters using ``get_space`` and then
initialize the underlying schedule using ``init_optimizer``
Parameters
----------
name: str
Name of a registered schedule
schedule: LRSchedule
Custom schedule, mutually exclusive with :param name
Examples
--------
.. code-block:: python
from olympus.optimizers import Optimizer
optimizer = Optimizer('sgd')
schedule = LRSchedule('exponential')
schedule.get_space()
# {'gamma': 'loguniform(0.97, 1)'}
schedule.init(optimizer, gamma=0.97)
Raises
------
RegisteredLRSchedulerNotFound
when using a name of an known schedule
MissingArgument:
if name nor schedule were not set
"""
def __init__(self, name=None, *, schedule=None, optimizer=None, **kwargs):
self._schedule = None
self._schedule_builder = None
self._optimizer = optimizer
self.hyper_parameters = HyperParameters(space={})
if schedule:
if isinstance(schedule, type):
self._schedule_builder = schedule
if hasattr(schedule, 'get_space'):
self.hyper_parameters.space = schedule.get_space()
else:
self._schedule = schedule
if hasattr(self._schedule, 'get_space'):
self.hyper_parameters.space = self._schedule.get_space()
elif name:
# load an olympus model
builder = registered_schedules.get(name)
if not builder:
raise RegisteredLRSchedulerNotFound(name)
self._schedule_builder = builder
if hasattr(self._schedule_builder, 'get_space'):
self.hyper_parameters.space = self._schedule_builder.get_space()
else:
raise MissingArgument('None or name needs to be set')
self.hyper_parameters.add_parameters(**kwargs)
def init(self, optimizer=None, override=False, **kwargs):
"""Initialize the LR schedule with the given hyper parameters"""
if self._schedule:
warning('LRSchedule is already set, use override=True to force re initialization')
if not override:
return self._schedule
if optimizer is None:
optimizer = self._optimizer
if optimizer is None:
raise MissingArgument('Missing optimizer argument!')
self.hyper_parameters.add_parameters(**kwargs)
self._schedule = self._schedule_builder(
optimizer,
**self.hyper_parameters.parameters(strict=True))
return self
def get_space(self):
"""Return the missing hyper parameters required to initialize the LR schedule"""
if self._schedule:
warning('LRSchedule is already set')
return self.hyper_parameters.missing_parameters()
def get_current_space(self):
"""Get currently defined parameter space"""
return self.hyper_parameters.parameters(strict=False)
@property
def defaults(self):
"""Return default hyper parameters"""
return self._schedule_builder.defaults()
@property
def lr_scheduler(self):
if not self._schedule:
self.init()
return self._schedule
def state_dict(self, destination=None, prefix='', keep_vars=False):
return self.lr_scheduler.state_dict()
def load_state_dict(self, state_dict, strict=True):
self.lr_scheduler.load_state_dict(state_dict)
def epoch(self, epoch, metrics=None):
"""Called after every epoch to update LR"""
self.lr_scheduler.epoch(epoch, metrics)
def step(self, step, metrics=None):
"""Called every step/batch to update LR"""
self.lr_scheduler.step(step, metrics)
def get_lr(self):
return self.lr_scheduler.get_lr()
class Model(nn.Module):
"""Olympus standardized Model interface
Parameters
----------
name: str
Name of a registered model
half: bool
Convert the network to half/fp16
model: Model
Custom model to use, mutually exclusive with :param name
Examples
--------
Model wrappers that provide a wide range of utility built-in.
Can instantiate common model directly
>>> model = Model('resnet18', input_size=(1, 28, 28), output_size=(10,))
Handles mixed precision conversion for you
>>> model = Model('resnet18', input_size=(1, 28, 28), output_size=(10,), half=True)
Handles weight initialization
>>> model = Model('resnet18', input_size=(1, 28, 28), output_size=(10,), weight_init='glorot_uniform')
Supports your custom model
>>> class MyModel(nn.Module):
... def __init__(self, input_size, output_size):
... self.main = nn.Linear(input_size[0], output_size[0])
...
... def forward(self, x):
... return self.main(x)
>>>
>>> model = Model(model=MyModel, input_size=(1, 28, 28), output_size=(10,))
Raises
------
RegisteredModelNotFound
when using a name of an known model
MissingArgument:
if name nor model were not set
"""
_dtype = torch.float32
_device = torch.device('cpu')
def __init__(self,
name=None,
*,
half=False,
model=None,
input_size=None,
output_size=None,
weight_init=default_init,
**kwargs):
super(Model, self).__init__()
# Save all the args that ware passed down so we can instantiate it again in standalone
self.replay_args = dict(name=name,
half=half,
model=model,
input_size=input_size,
output_size=output_size,
weight_init=weight_init,
kwargs=kwargs)
self.transform = lambda x: try_convert(x, self.device, self.dtype)
self.half = half
self._model = None
# Track defined hyper parameters
self.hyper_parameters = HyperParameters(space=dict())
# If init is set then we can add its hyper parameters
self.weight_init = weight_init
if weight_init is not None:
if isinstance(weight_init, str):
self.weight_init = Initializer(weight_init)
# replace weight init by its own hyper parameters
space = self.weight_init.get_space()
if space:
self.hyper_parameters.space.update(dict(initializer=space))
# Make a Lazy Model that will be initialized once all the hyper parameters are set
if model:
if hasattr(model, 'get_space'):
self.hyper_parameters.space.update(model.get_space())
if isinstance(model, type):
self.model_builder = LazyCall(model,
input_size=input_size,
output_size=output_size)
else:
self.model_builder = LazyCall(lambda *args, **kwargs: model)
elif name:
# load an olympus model
model_fun = registered_models.get(name)
if not model_fun:
raise RegisteredModelNotFound(name)
self.model_builder = LazyCall(model_fun,
input_size=input_size,
output_size=output_size)
if hasattr(model_fun, 'get_space'):
self.hyper_parameters.space.update(model_fun.get_space())
else:
raise MissingArgument('Model or Name need to be set')
# Any Additional parameters set Hyper parameters
self.other_params = self.hyper_parameters.add_parameters(strict=False,
**kwargs)
@property
def dtype(self):
return self._dtype
@property
def device(self):
return self._device
def get_space(self):
"""Return hyper parameter space"""
return self.hyper_parameters.missing_parameters()
def get_current_space(self):
"""Get currently defined parameter space"""
return self.hyper_parameters.parameters(strict=False)
def init(self, override=False, **model_hyperparams):
others = self.hyper_parameters.add_parameters(strict=False,
**model_hyperparams)
self.other_params.update(others)
params = self.hyper_parameters.parameters(strict=True)
initializer = params.pop('initializer', {})
if isinstance(initializer, dict):
self.weight_init.init(**initializer)
self._model = self.model_builder.invoke(**self.other_params, **params)
self.weight_init(self._model)
if self.half:
self._model = network_to_half(self._model)
# Register module so we can use all the parent methods
self.add_module('_model', self._model)
self.replay_args['kwargs'].update(model_hyperparams)
return self
@property
def model(self):
if not self._model:
self.init()
return self._model
def forward(self, *input, **kwargs):
return self.model(self.transform(input[0]), *input[1:], **kwargs)
def __call__(self, *args, **kwargs):
return super(Model, self).__call__(*args, **kwargs)
def state_dict(self, destination=None, prefix='', keep_vars=False):
destination = {
'model': self.model.state_dict(None, prefix, keep_vars),
'half': self.half,
'replay': self.replay_args,
'types': {
'model': type(self.model)
}
}
return destination
@staticmethod
def from_state(state):
kwargs = state.get('replay')
kwargs.update(kwargs.pop('kwargs', dict()))
m = Model(**kwargs)
m.init()
m.load_state_dict(state)
return m
def load_state_dict(self, state_dict, strict=True):
self.half = state_dict['half']
self.model.load_state_dict(state_dict['model'], strict=strict)
def parameters(self, recurse: bool = True):
return self.model.parameters(recurse)
def to(self, *args, **kwargs):
self._device, self._dtype, *_ = torch._C._nn._parse_to(*args, **kwargs)
super(Model, self).to(*args, **kwargs)
return self
def act(self, *args, **kwargs):
return self.model.act(*args, **kwargs)
def critic(self, *args, **kwargs):
return self.model.critic(*args, **kwargs)