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