Skip to content

Commit

Permalink
Add: Support custom optimization algorithms
Browse files Browse the repository at this point in the history
  • Loading branch information
@andreArtelt
andreArtelt committed Aug 1, 2019
1 parent 16a3bf0 commit 1e803de
Show file tree
Hide file tree
Showing 18 changed files with 635 additions and 239 deletions.
89 changes: 57 additions & 32 deletions ceml/backend/tensorflow/optimizer/optimizer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,40 +4,62 @@
import inspect
from ..layer import create_tensor, create_mutable_tensor
from ....optim import Optimizer
from ....optim import desc_to_optim as desc_to_optim_scipy
from ....optim import prepare_optim as desc_to_optim_scipy


class TfOptimizer(Optimizer):
"""Wrapper for a tensorflow optimization algorithm.
The :class:`TfOptimizer` provides an interface for wrapping an arbitrary tensorflow optimization algorithm (see :class:`tf.train.Optimizer`) and minimizing a given loss function.
Parameters
----------
model : `callable` or instance of :class:`tf.keras.Model`
The model that is to be used.
loss : instance of :class:`ceml.backend.tensorflow.costfunctions.RegularizedCost`
The loss that has to be minimized.
x : `numpy.ndarray`
The starting value of `x` - usually this is the original input whose prediction has to be explained..
optim : instance of :class:`tf.train.Optimizer`
Optimizer for minimizing the loss.
tol : `float`, optional
Tolerance for termination.
The default is 0.0
max_iter : `int`, optional
Maximum number of iterations.
The default is 1.
grad_mask : `numpy.array`, optional
Mask that is multiplied element wise on top of the gradient - can be used to hold some dimensions constant.
If `grad_mask` is None, no gradient mask is used.
The default is None.
"""
def __init__(self, model, loss, x, optim, tol=None, max_iter=1, grad_mask=None):
def __init__(self):
self.model = None
self.loss = None
self.tol = None
self.max_iter = None

self.x = None

self.optim = None

self.grad_mask = None

super(TfOptimizer, self).__init__()

def init(self, model, loss, x, optim, tol=None, max_iter=1, grad_mask=None):
"""
Initializes all parameters.
Parameters
----------
model : `callable` or instance of :class:`tf.keras.Model`
The model that is to be used.
loss : instance of :class:`ceml.backend.tensorflow.costfunctions.RegularizedCost`
The loss that has to be minimized.
x : `numpy.ndarray`
The starting value of `x` - usually this is the original input whose prediction has to be explained..
optim : instance of :class:`tf.train.Optimizer`
Optimizer for minimizing the loss.
tol : `float`, optional
Tolerance for termination.
The default is 0.0
max_iter : `int`, optional
Maximum number of iterations.
The default is 1.
grad_mask : `numpy.array`, optional
Mask that is multiplied element wise on top of the gradient - can be used to hold some dimensions constant.
If `grad_mask` is None, no gradient mask is used.
The default is None.
Raises
------
TypeError
If the type of `loss` or `model` is not correct.
"""
if not callable(model):
raise TypeError("model must be callable")
if not callable(loss):
Expand All @@ -56,9 +78,10 @@ def __init__(self, model, loss, x, optim, tol=None, max_iter=1, grad_mask=None):
self.grad_mask = create_tensor(grad_mask)
else:
self.grad_mask = create_tensor(np.ones_like(x))

super(TfOptimizer, self).__init__()

def is_grad_based(self):
return True

def __loss_grad(self):
with tf.GradientTape(watch_accessed_variables=False) as tape:
tape.watch(self.x)
Expand Down Expand Up @@ -86,10 +109,12 @@ def __call__(self):
return self.minimize()


def desc_to_optim(optimizer, loss, loss_npy, loss_grad_npy, x_orig, model, tol, max_iter, grad_mask):
if isinstance(optimizer, str):
def prepare_optim(optimizer, loss, loss_npy, loss_grad_npy, x_orig, model, tol, max_iter, grad_mask):
if isinstance(optimizer, str) or isinstance(optimizer, Optimizer):
return desc_to_optim_scipy(optimizer, loss_npy, x_orig, loss_grad_npy, tol, max_iter)
elif isinstance(optimizer, tf.compat.v1.train.Optimizer):
return TfOptimizer(model, loss, x_orig, optimizer, tol, max_iter, grad_mask)
optim = TfOptimizer()
optim.init(model, loss, x_orig, optimizer, tol, max_iter, grad_mask)
return optim
else:
raise TypeError("Invalid type of argument 'optimizer'.\n'optimizer' has to be a string or the name of a class that is derived from tf.train.Optimizer")
122 changes: 72 additions & 50 deletions ceml/backend/torch/optimizer/optimizer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,59 +4,78 @@
import inspect
from ..layer import create_tensor
from ....optim import Optimizer
from ....optim import desc_to_optim as desc_to_optim_scipy
from ....optim import prepare_optim as desc_to_optim_scipy


class TorchOptimizer(Optimizer):
"""Wrapper for a PyTorch optimization algorithm.
The :class:`TorchOptimizer` provides an interface for wrapping an arbitrary PyTorch optimization algorithm (see :class:`torch.optim`) and minimizing a given loss function.
"""
def __init__(self):
self.model = None
self.loss = None
self.tol = None
self.max_iter = None
self.device = None

self.x = None

self.optim = None
self.lr_scheduler = None

self.grad_mask = None

Parameters
----------
model : instance of :class:`torch.nn.Module`
The model that is to be used.
loss : instance of :class:`ceml.backend.torch.costfunctions.RegularizedCost`
The loss that has to be minimized.
x : `numpy.ndarray`
The starting value of `x` - usually this is the original input whose prediction has to be explained..
optim : instance of `torch.optim.Optimizer`
Optimizer for minimizing the loss.
optim_args : `dict`
Arguments of the optimization algorithm (e.g. learning rate, momentum, ...)
lr_scheduler : Learning rate scheduler (see :class:`torch.optim.lr_scheduler`)
Learning rate scheduler (see :class:`torch.optim.lr_scheduler`).
The default is None.
lr_scheduler_args : `dict`, optional
Arguments of the learning rate scheduler.
The default is None.
tol : `float`, optional
Tolerance for termination.
The default is 0.0
max_iter : `int`, optional
Maximum number of iterations.
The default is 1.
grad_mask : `numpy.array`, optional
Mask that is multiplied element wise on top of the gradient - can be used to hold some dimensions constant.
If `grad_mask` is None, no gradient mask is used.
The default is None.
device : :class:`torch.device`
Specifies the hardware device (e.g. cpu or gpu) we are working on.
The default is `torch.device("cpu")`.
super(TorchOptimizer, self).__init__()

Raises
------
TypeError
If the parameters loss or model do not have the correct type.
"""
def __init__(self, model, loss, x, optim, optim_args, lr_scheduler=None, lr_scheduler_args=None, tol=None, max_iter=1, grad_mask=None, device=torch.device("cpu")):
def init(self, model, loss, x, optim, optim_args, lr_scheduler=None, lr_scheduler_args=None, tol=None, max_iter=1, grad_mask=None, device=torch.device("cpu")):
"""
Initializes all parameters.
Parameters
----------
model : instance of :class:`torch.nn.Module`
The model that is to be used.
loss : instance of :class:`ceml.backend.torch.costfunctions.RegularizedCost`
The loss that has to be minimized.
x : `numpy.ndarray`
The starting value of `x` - usually this is the original input whose prediction has to be explained..
optim : instance of `torch.optim.Optimizer`
Optimizer for minimizing the loss.
optim_args : `dict`
Arguments of the optimization algorithm (e.g. learning rate, momentum, ...)
lr_scheduler : Learning rate scheduler (see :class:`torch.optim.lr_scheduler`)
Learning rate scheduler (see :class:`torch.optim.lr_scheduler`).
The default is None.
lr_scheduler_args : `dict`, optional
Arguments of the learning rate scheduler.
The default is None.
tol : `float`, optional
Tolerance for termination.
The default is 0.0
max_iter : `int`, optional
Maximum number of iterations.
The default is 1.
grad_mask : `numpy.array`, optional
Mask that is multiplied element wise on top of the gradient - can be used to hold some dimensions constant.
If `grad_mask` is None, no gradient mask is used.
The default is None.
device : :class:`torch.device`
Specifies the hardware device (e.g. cpu or gpu) we are working on.
The default is `torch.device("cpu")`.
Raises
------
TypeError
If the type of `loss` or `model` is not correct.
"""
if not isinstance(model, torch.nn.Module):
raise TypeError(f"model must be an instance of torch.nn.Module not of {type(model)}")
if not callable(loss):
Expand All @@ -81,8 +100,9 @@ def __init__(self, model, loss, x, optim, optim_args, lr_scheduler=None, lr_sche
else:
self.grad_mask = self.x.new_ones(self.x.shape)

super(TorchOptimizer, self).__init__()

def is_grad_based(self):
return True

def minimize(self):
old_loss = float("inf")
for _ in range(self.max_iter):
Expand All @@ -107,10 +127,12 @@ def __call__(self):
return self.minimize()


def desc_to_optim(optimizer, optimizer_args, lr_scheduler, lr_scheduler_args, loss, loss_npy, loss_grad_npy, x_orig, model, tol, max_iter, grad_mask, device):
if isinstance(optimizer, str):
def prepare_optim(optimizer, optimizer_args, lr_scheduler, lr_scheduler_args, loss, loss_npy, loss_grad_npy, x_orig, model, tol, max_iter, grad_mask, device):
if isinstance(optimizer, str) or isinstance(optimizer, Optimizer):
return desc_to_optim_scipy(optimizer, loss_npy, x_orig, loss_grad_npy, tol, max_iter)
elif inspect.isclass(optimizer) == True:
return TorchOptimizer(model, loss, x_orig, optimizer, optimizer_args, lr_scheduler, lr_scheduler_args, tol, max_iter, grad_mask, device)
optim = TorchOptimizer()
optim.init(model, loss, x_orig, optimizer, optimizer_args, lr_scheduler, lr_scheduler_args, tol, max_iter, grad_mask, device)
return optim
else:
raise TypeError("Invalid type of argument 'optimizer'.\n'optimizer' has to be a string or the name of a class that is derived from torch.optim.Optimizer")

0 comments on commit 1e803de

Please sign in to comment.