def call(self, *args, **kwargs) -> Any:
"""
Makes a module callable and contains the forward pass of your model.
This should be pure computation and not allocate any weights.
Allocating weights should be done in the `build` function.
This function gets called by `__call__` and itself passes all calls to `_call_pytorch` and `_call_tf`.
Furthermore, it takes care of unwrapping the tensors into native tensors before calling and wrapping them again after calling.
This allows the native functions `_call_pytorch` and `_call_tf` to be pure pytorch or tensorflow code.
All subclasses must implement `_call_pytorch` and `_call_tf`.
You should call this module in the following style (this ensures the module is build on first run):
```
module = MyModule()
result = module(*args, **kwargs)
```
Parameters:
:param *args: You can specify any parameters you want.
:param **kwargs: You can specify any named parameters you want.
"""
args = self._wrapper.unwrap(args)
kwargs = self._wrapper.unwrap(kwargs)
if is_backend(PYTORCH_BACKEND):
results = self._call_pytorch(*args, **kwargs)
elif is_backend(TF_BACKEND):
results = self._call_tf(*args, **kwargs)
else:
raise RuntimeError("Unknown Backend: {}".format(get_backend()))
results = self._wrapper.wrap(results)
return results
def build(self, *args, **kwargs) -> None:
"""
Build the model, this function automatically calls the native build with the tensors unwrapped.
This function gets called by `__call__` and itself passes all calls to `_build_pytorch` and `_build_tf`.
Furthermore, it takes care of unwrapping the tensors into native tensors before calling and wrapping them again after calling.
This allows the native functions `_build_pytorch` and `_build_tf` to be pure pytorch or tensorflow code.
All subclasses must implement `_build_pytorch` and `_build_tf`.
You should never call the build function directly. Call this module in the following style (this ensures the module is build on first run):
```
module = MyModule()
result = module(*args, **kwargs) # <- Build gets called internally here.
```
Parameters:
:param *args: You must specify the exact same parameters as for your call.
:param **kwargs: You must specify the exact same parameters as for your call.
"""
args = self._wrapper.unwrap(args)
kwargs = self._wrapper.unwrap(kwargs)
if is_backend(PYTORCH_BACKEND):
self._build_pytorch(*args, **kwargs)
elif is_backend(TF_BACKEND):
self._build_tf(*args, **kwargs)
else:
raise RuntimeError("Unknown Backend: {}".format(get_backend()))
def __exit__(self, type, value, traceback):
_device_stack.pop()
if babilim.is_backend(TF_BACKEND):
self.native_device.__exit__()
self.native_device = None
elif babilim.is_backend(PYTORCH_BACKEND):
pass
else:
raise RuntimeError(
"No implementation for this backend was found. (backend={})".
format(babilim.get_backend()))
def __enter__(self):
_device_stack.append(self.name)
if babilim.is_backend(TF_BACKEND):
import tensorflow as tf
self.native_device = tf.device(get_current_device_native_format())
self.native_device.__enter__()
elif babilim.is_backend(PYTORCH_BACKEND):
pass
else:
raise RuntimeError(
"No implementation for this backend was found. (backend={})".
format(babilim.get_backend()))
return self
def log(tensor: ITensor) -> ITensor:
"""
Compute the logarithm of a tensor.
:param tensor: The tensor of which the log should be computed.
:return: The log tensor.
"""
if is_backend(PYTORCH_BACKEND):
from babilim.core.tmath.pytorch import log as _log
return _log(tensor)
elif is_backend(TF_BACKEND):
from babilim.core.tmath.tf import log as _log
return _log(tensor)
def pow(tensor: ITensor, exponent: ITensor) -> ITensor:
"""
Compute the power of a tensor.
:param tensor: The tensor of which the pow should be computed.
:param exponent: The exponent to take the power with.
:return: The pow tensor.
"""
if is_backend(PYTORCH_BACKEND):
from babilim.core.tmath.pytorch import pow as _pow
return _pow(tensor, exponent)
elif is_backend(TF_BACKEND):
from babilim.core.tmath.tf import pow as _pow
return _pow(tensor, exponent)
def to_dataloader(self) -> Dataloader:
"""
Converts the dataset into a babilim.data.Dataloader.
:return: Returns a babilim dataloader object usable with the trainers.
"""
data = None
if babilim.is_backend(babilim.PYTORCH_BACKEND):
data = self.to_pytorch()
elif babilim.is_backend(babilim.TF_BACKEND):
data = self.to_keras()
else:
raise NotImplementedError(
"Other backends than pytorch and tf2 are not implemented.")
return Dataloader(data, self)
def clip(tensor: ITensor, min, max) -> ITensor:
"""
Clip a tensor to a value range.
:param tensor: The tensor which should be clipped.
:param min: The minimum value to clip to.
:param max: The maximum value to clip to.
:return: The clipped tensor.
"""
if is_backend(PYTORCH_BACKEND):
from babilim.core.tmath.pytorch import clip as _clip
return _clip(tensor, min, max)
elif is_backend(TF_BACKEND):
from babilim.core.tmath.tf import clip as _clip
return _clip(tensor, min, max)
def load_state(checkpoint_path: str, native_format: bool = False) -> Dict:
"""
Load the state from a checkpoint.
:param checkpoint_path: The path to the file in which the checkpoint is stored.
:param native_format: (Optional) If the checkpoint should use the backend specific native format. (default: False)
:return: A dict containing the states.
"""
if native_format:
if is_backend(PYTORCH_BACKEND):
import torch
return torch.load(checkpoint_path, map_location='cpu')
else:
raise NotImplementedError()
else:
data = np.load(checkpoint_path, allow_pickle=False)
out = {}
prefixes = list(set([key.split("/")[0] for key in list(data.keys())]))
for prefix in prefixes:
if prefix in data: # primitive types
out[prefix] = data[prefix]
else: # dict types
tmp = {"{}".format("/".join(k.split("/")[1:])): data[k] for k in data if k.startswith(prefix)}
out[prefix] = tmp
return out
def save_state(data, checkpoint_path, native_format=False):
"""
Save the state to a checkpoint.
:param data: A dict containing the states.
:param checkpoint_path: The path to the file in which the checkpoint shall be stored.
:param native_format: (Optional) If the checkpoint should use the backend specific native format. (default: False)
"""
if native_format:
if is_backend(PYTORCH_BACKEND):
import torch
return torch.save(data, checkpoint_path)
else:
raise NotImplementedError()
else:
out = {}
for key, value in data.items():
if isinstance(value, dict):
tmp = {"{}/{}".format(key, k): value[k] for k in value}
out.update(tmp)
elif any(isinstance(value, t) for t in [int, str, float, list]):
out[key] = value
else:
raise RuntimeError("The type ({}) of {} is not allowed!".format(type(value), key))
np.savez_compressed(checkpoint_path, **out)
def image_grid_wrap(data: np.ndarray) -> np.ndarray:
"""
Prepares 2D grid information to be used in your neural network.
This is required for all data that should be usable by a 2D Convolution.
If your data has shape [H, W] it gets transformed to [H, W, 1] automatically.
For pytorch the data gets further reordered in a channel first order [C, H, W].
:param data: The data that should be prepared.
:return: A numpy ndarray with the prepared image/grid data.
"""
if data.ndim != 3 and data.ndim != 2:
raise ValueError(
"Wrong dimensionality of the data. Must be 2/3 dimensional but is {} dimensional."
.format(data.ndim))
# Add a dimension at the end.
if data.ndim == 2:
data = data[:, :, None]
# Transpose data for numpy.
if babilim.is_backend(babilim.PYTORCH_BACKEND):
return data.transpose((2, 0, 1))
else:
return data
def get_current_device_native_format() -> str:
"""
Get a string specifying the currently selected default device in the backend specific native format.
When you manually assign a device, you should always use this device.
"""
name = _device_stack[-1]
if babilim.is_backend(TF_BACKEND):
return "/" + name
elif babilim.is_backend(PYTORCH_BACKEND):
import torch
if torch.cuda.is_available():
return name.replace("gpu", "cuda")
else:
return "cpu"
else:
raise RuntimeError(
"No implementation for this backend was found. (backend={})".
format(babilim.get_backend()))
def _register_params(self, module):
"""
Allows registration of the parameters with a native module.
This makes the parameters of a babilim modules available to the native modules.
When using a babilim modules in a native modules, use this function and pass the native module as a parameter.
This function works by adding all trainable_variables to the module you pass.
Warning: You need to build the babilim modules before calling this function. Building can be done by calling for example.
Here is a pytorch example:
```python
import torch
from torch.nn import Module
from babilim.modules import Linear
class MyModule(Module):
def __init__(self):
super().__init__()
self.linear = Linear(10)
def forward(self, features):
result = self.linear(features)
self.linear.register_params(self)
return result
```
:param module: The native module on which parameters of this modules should be registered.
"""
if babilim.is_backend(PYTORCH_BACKEND):
from torch.nn import Module
if isinstance(module, Module):
myname = "_error_"
for var in module.__dict__:
if module.__dict__[var] == self:
myname = var
if isinstance(module.__dict__[var],
list) and self in module.__dict__[var]:
myname = "{}/{}".format(
var, module.__dict__[var].index(self))
# Register self as pytorch module.
module._modules[myname] = self
for name, param in self.named_variables.items():
if param.trainable:
module.register_parameter(myname + name, param.native)
else:
module.register_buffer(myname + name, param.native)
else:
if babilim.core.logging.DEBUG_VERBOSITY:
_warn_once(
"babilim.model.module.Module:_register_params Not implemented for tf2 but I think it is not required."
)
def save(self, name: str) -> None:
chkpt_extension = ".npz"
if self.native_format:
if babilim.is_backend(babilim.TF2_BACKEND):
chkpt_extension = ""
elif babilim.is_backend(babilim.PYTORCH_BACKEND):
chkpt_extension = ".pth"
checkpoint_sub_path = os.path.join(
"checkpoints", "{}{}".format(name, chkpt_extension))
checkpoint_path = os.path.join(get_log_path(), checkpoint_sub_path)
save_state(
{
"epoch": self.epoch,
"model": self.model.state_dict(),
"optimizer": self.optimizer.state_dict(),
"loss": self.loss.state_dict()
},
checkpoint_path,
native_format=self.native_format)
info("Saved Checkoint: {}".format(checkpoint_sub_path))
def state_dict(self) -> Dict:
"""
Get the state of the object as a state dict (usable for checkpoints).
:return: A dictionary containing the state of the object.
"""
state = {}
for name, var in self.named_variables.items():
if babilim.is_backend(babilim.TF_BACKEND):
state[name] = var.numpy()
else:
state[name] = var.numpy().T
return state
def loss(self, y_pred: ITensor, y_true: ITensor) -> ITensor:
"""
Compute the sparse cross entropy assuming y_pred to be logits.
:param y_pred: The predictions of the network. Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
:param y_true: The desired outputs of the network (labels). Either a NamedTuple pointing at ITensors or a Dict or Tuple of ITensors.
"""
if babilim.is_backend(babilim.PYTORCH_BACKEND):
return Tensor(data=self.loss_fun(y_pred.native, y_true.native),
trainable=True)
else:
return Tensor(data=self.loss_fun(labels=y_true.native,
logits=y_pred.native),
trainable=True)
def __init__(self, reduction: str = "mean"):
"""
Compute a binary cross entropy.
This means that the preds are logits and the targets are a binary (1 or 0) tensor of same shape as logits.
:param reduction: Specifies the reduction to apply to the output: `'none'` | `'mean'` | `'sum'`. `'none'`: no reduction will be applied, `'mean'`: the sum of the output will be divided by the number of elements in the output, `'sum'`: the output will be summed. Default: 'mean'.
"""
super().__init__(reduction=reduction)
if babilim.is_backend(babilim.PYTORCH_BACKEND):
from torch.nn import BCEWithLogitsLoss
self.loss_fun = BCEWithLogitsLoss(reduction="none")
else:
from tensorflow.nn import sigmoid_cross_entropy_with_logits
self.loss_fun = sigmoid_cross_entropy_with_logits
def __init__(self, reduction: str = "mean"):
"""
Compute a sparse cross entropy.
This means that the preds are logits and the targets are not one hot encoded.
:param reduction: Specifies the reduction to apply to the output: `'none'` | `'mean'` | `'sum'`. `'none'`: no reduction will be applied, `'mean'`: the sum of the output will be divided by the number of elements in the output, `'sum'`: the output will be summed. Default: 'mean'.
"""
super().__init__(reduction=reduction)
if babilim.is_backend(babilim.PYTORCH_BACKEND):
from torch.nn import CrossEntropyLoss
self.loss_fun = CrossEntropyLoss(reduction="none")
else:
from tensorflow.nn import sparse_softmax_cross_entropy_with_logits
self.loss_fun = sparse_softmax_cross_entropy_with_logits
def image_grid_unwrap(data: np.ndarray) -> np.ndarray:
"""
For pytorch the data gets reordered in a channel last order [H, W, C].
:param data: The data that should be unwrapped.
:return: A numpy ndarray with the unwrapped image/grid data.
"""
if data.ndim != 3:
raise ValueError(
"Wrong dimensionality of the data. Must be 3 dimensional but is {} dimensional."
.format(data.ndim))
# Transpose data for numpy.
if babilim.is_backend(babilim.PYTORCH_BACKEND):
return data.transpose((1, 2, 0))
else:
return data
def load_state_dict(self, state_dict: Dict) -> None:
"""
Load the state of the object from a state dict.
Handy when loading checkpoints.
:param state_dict: A dictionary containing the state of the object.
"""
for name, var in self.named_variables.items():
if name in state_dict:
if babilim.is_backend(babilim.TF_BACKEND):
var.assign(state_dict[name])
else:
var.assign(state_dict[name].T)
if DEBUG_VERBOSITY:
info(" Loaded: {}".format(name))
else:
warn(" Variable {} not in checkpoint.".format(name))
def _auto_device(self):
if babilim.is_backend(babilim.PYTORCH_BACKEND):
import torch
self.native_loss = self.native_loss.to(torch.device(self.device))
return self