notebooks/custom_backend.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How To: Integrate your own Computation Backend\n",
    "\n",
    "*Author: Justus Schock*\n",
    "\n",
    "*Date: 15.05.2019*\n",
    "\n",
    "This howto will take you on a trip through the `delira` internals, while we will see, how to add a custom computation backend on the examplaric case of the `torch.jit` or `TorchScript` backend\n",
    "\n",
    "## Model Definitions\n",
    "In order to implement a network, we will first have to define the network itself. In `delira` there is a single backend-specific implementation of an abstract network class for each of the backends. These interface classes are all based on the `AbstractNetwork`-class, defining the major API.\n",
    "\n",
    "So let's start having a look at this class to see, what we will have to implement for our own backend.\n",
    "\n",
    "Of course we will have to implement an `__init__` defining our class. The `__init__` of `AbstractNetwork` (which should be called during our the `__init__` of our baseclass) accepts a number of kwargs and simply registers them to be `init_kwargs`, so there is nothing we have to take care of.\n",
    "\n",
    "The next function to inspect is the `__call__` function, which makes the class callable and the docstrings indicate, that it should take care of our model's forward-pass.\n",
    "\n",
    "After the `__call__` we now have the `closure` function, which defines a single training step (including, but not limited to, forward-pass, calculation of losses and train-metrics, backward-pass and optimization).\n",
    "\n",
    "The last method to implement is the `prepare_batch` function which converts the input to a suitable format and the correct data-type and device.\n",
    "\n",
    "### TorchScript Limitations\n",
    "Since we want to implement an abstract network class for this specific backend, we should have a look on how to generally implement models in this backend.\n",
    "\n",
    "According the the [PyTorch docs](https://pytorch.org/docs/stable/jit.html) this works as follows:\n",
    "\n",
    "> You can write TorchScript code directly using Python syntax. You do this using the `torch.jit.script` decorator (for functions) or `torch.jit.script_method` decorator (for methods) on subclasses of `ScriptModule`. With this decorator the body of the annotated function is directly translated into TorchScript. TorchScript itself is a subset of the Python language, so not all features in Python work, but we provide enough functionality to compute on tensors and do control-dependent operations.\n",
    "\n",
    "Since our use-case is to implement the interface class for networks, we want to use the way of subclassing `torch.jit.ScriptModule`, implement it's `forward` and use the `torch.jit.script_method` decorator on it.\n",
    "\n",
    "The example given in the very same docs for this case is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.4997, 0.2955, 0.1588, 0.1873, 0.4753], grad_fn=<MvBackward>)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "class MyScriptModule(torch.jit.ScriptModule):\n",
    "    def __init__(self, N, M):\n",
    "        super().__init__()\n",
    "        self.weight = torch.nn.Parameter(torch.rand(N, M))\n",
    "\n",
    "    @torch.jit.script_method\n",
    "    def forward(self, input):\n",
    "        return self.weight.mv(input)\n",
    "    \n",
    "my_script_module = MyScriptModule(5, 3)\n",
    "input_tensor = torch.rand(3)\n",
    "my_script_module(input_tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Merging TorchScript into our Abstract Class\n",
    "\n",
    "This little example gives us a few things, we have to do for a successful definition of our base class:\n",
    "\n",
    "**1.)** Our class has to subclass both, the `AbstractNetwork` and the `torch.jit.ScriptModule` classes.\n",
    "\n",
    "**2.)** We need to implement a `forward` method, which takes care of the forward-pass (as it's name indicates).\n",
    "\n",
    "**3.)** We don't have to take care of the backward-pass (thanks to `PyTorch`'s and `TorchScript`'s AutoGrad (which is a framework for automatic differentiation).\n",
    "\n",
    "**4.)** Since `torch.jit.ScriptModule` is callable (seen in the example), it already implements a `__call__` method and we may simply use this one.\n",
    "\n",
    "**5.)** The `closure` is completely network-dependent and thus has to remain an abstract method here.\n",
    "\n",
    "**6.)** The `prepare_batch` function also depends on the combination of network, inputs and loss functions to use, but we can at least give a prototype of such an function, which handles the devices correctly and converts everything to `float`\n",
    "\n",
    "\n",
    "### Actual Implementation\n",
    "\n",
    "Now, let's start with the actual implementation and do one function by another and keep the things in mind, we just discovered.\n",
    "\n",
    "#### Class Signature and `__init__`-Method\n",
    "To subclass both networks, we cannot use the simple `super().__init__` approach, because we have to init both parent classes, so we do \n",
    "\n",
    "\n",
    "```python\n",
    "\n",
    "    class AbstractTorchScriptNetwork(AbstractNetwork, torch.jit.ScriptModule):\n",
    "\n",
    "        @abc.abstractmethod\n",
    "        def __init__(self, optimize=True, **kwargs):\n",
    "            \"\"\"\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            optimize : bool\n",
    "                whether to optimize the network graph or not; default: True\n",
    "            **kwargs :\n",
    "                additional keyword arguments (passed to :class:`AbstractNetwork`)\n",
    "            \"\"\"\n",
    "            torch.jit.ScriptModule.__init__(self, optimize=optimize)\n",
    "            AbstractNetwork.__init__(self, **kwargs)\n",
    "            \n",
    "```\n",
    "instead. This ensures all parent classes to be initialized correctly.\n",
    "\n",
    "#### `__call__`-Method\n",
    "As mentioned above, the `__call__` method is very easy to implement, because we can simply use the implementation of our `TorchScript` base class like this:\n",
    "\n",
    "```python\n",
    "\n",
    "    def __call__(self, *args, **kwargs):\n",
    "        \"\"\"\n",
    "        Calls Forward method\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        *args :\n",
    "            positional arguments (passed to `forward`)\n",
    "        **kwargs :\n",
    "            keyword arguments (passed to `forward`)\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        Any\n",
    "            result: module results of arbitrary type and number\n",
    "\n",
    "        \"\"\"\n",
    "        return torch.jit.ScriptModule.__call__(self, *args, **kwargs)\n",
    "        \n",
    "```\n",
    "\n",
    "This also ensures, that we can pass an arbitrary number or positional and keyword arguments of arbitrary types to it (which are all passed to the `forward`-function). The advantage over directly calling the `forward` method here, is that the `ScriptModule.__call__` already does the handling of [forward-pre-hooks](https://pytorch.org/docs/stable/nn.html#torch.nn.Module.register_forward_pre_hook), [forward-hooks](https://pytorch.org/docs/stable/nn.html#torch.nn.Module.register_forward_hook) and [backward-hooks](https://pytorch.org/docs/stable/nn.html#torch.nn.Module.register_backward_hook).\n",
    "\n",
    "#### `closure`-Method\n",
    "Since this method is highly model-dependant, we just don't implement it, which forces the user to implement it (since it is marked as an `abstractmethod` in `AbstractExperiment`).\n",
    "\n",
    "#### `prepare_batch`-Method\n",
    "The above mentioned prototype of pushing everything to the correct device and convert it to float looks like this:\n",
    "\n",
    "```python\n",
    "\n",
    "    @staticmethod\n",
    "    def prepare_batch(batch: dict, input_device, output_device):\n",
    "        \"\"\"\n",
    "        Helper Function to prepare Network Inputs and Labels (convert them to\n",
    "        correct type and shape and push them to correct devices)\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        batch : dict\n",
    "            dictionary containing all the data\n",
    "        input_device : torch.device\n",
    "            device for network inputs\n",
    "        output_device : torch.device\n",
    "            device for network outputs\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        dict\n",
    "            dictionary containing data in correct type and shape and on correct\n",
    "            device\n",
    "\n",
    "        \"\"\"\n",
    "        return_dict = {\"data\": torch.from_numpy(batch.pop(\"data\")).to(\n",
    "            input_device).to(torch.float)}\n",
    "\n",
    "        for key, vals in batch.items():\n",
    "            return_dict[key] = torch.from_numpy(vals).to(output_device).to(\n",
    "                torch.float)\n",
    "\n",
    "        return return_dict\n",
    "\n",
    "```\n",
    "\n",
    "Since we don't want to use any of the model's attributes here (and for conformity with the `AbstractNetwork` class), this method is defined as `staticmethod`, meaning it is class-bound, not instance-bound. The `closure` method has to be a `staticmethod` too.\n",
    "\n",
    "\n",
    "#### `forward`-Method\n",
    "The only thing left now, is the `forward` method, which is internally called by `ScriptModule.__call__`. The bad news is: We currently can't implement it. Subclassing a `ScriptModule` to overwrite a function decorated with `torch.jit.script_method` is not (yet) supported, but will be soon, once [this PR](https://github.com/pytorch/pytorch/pull/20503) is merged and released.\n",
    "\n",
    "For now: you simply have to implement this method in your own network despite the missing of an abstract interface-method.\n",
    "\n",
    "#### Putting it all together\n",
    "If we combine all the function implementations to one class, it looks like this:\n",
    "\n",
    "```python\n",
    "\n",
    "    class AbstractTorchScriptNetwork(AbstractNetwork, torch.jit.ScriptModule):\n",
    "\n",
    "        \"\"\"\n",
    "        Abstract Interface Class for TorchScript Networks. For more information\n",
    "        have a look at https://pytorch.org/docs/stable/jit.html#torchscript\n",
    "\n",
    "        Warnings\n",
    "        --------\n",
    "        In addition to the here defined API, a forward function must be\n",
    "        implemented and decorated with ``@torch.jit.script_method``\n",
    "\n",
    "        \"\"\"\n",
    "        @abc.abstractmethod\n",
    "        def __init__(self, optimize=True, **kwargs):\n",
    "            \"\"\"\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            optimize : bool\n",
    "                whether to optimize the network graph or not; default: True\n",
    "            **kwargs :\n",
    "                additional keyword arguments (passed to :class:`AbstractNetwork`)\n",
    "            \"\"\"\n",
    "            torch.jit.ScriptModule.__init__(self, optimize=optimize)\n",
    "            AbstractNetwork.__init__(self, **kwargs)\n",
    "\n",
    "        def __call__(self, *args, **kwargs):\n",
    "            \"\"\"\n",
    "            Calls Forward method\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            *args :\n",
    "                positional arguments (passed to `forward`)\n",
    "            **kwargs :\n",
    "                keyword arguments (passed to `forward`)\n",
    "\n",
    "            Returns\n",
    "            -------\n",
    "            Any\n",
    "                result: module results of arbitrary type and number\n",
    "\n",
    "            \"\"\"\n",
    "            return torch.jit.ScriptModule.__call__(self, *args, **kwargs)\n",
    "\n",
    "        @staticmethod\n",
    "        def prepare_batch(batch: dict, input_device, output_device):\n",
    "            \"\"\"\n",
    "            Helper Function to prepare Network Inputs and Labels (convert them to\n",
    "            correct type and shape and push them to correct devices)\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            batch : dict\n",
    "                dictionary containing all the data\n",
    "            input_device : torch.device\n",
    "                device for network inputs\n",
    "            output_device : torch.device\n",
    "                device for network outputs\n",
    "\n",
    "            Returns\n",
    "            -------\n",
    "            dict\n",
    "                dictionary containing data in correct type and shape and on correct\n",
    "                device\n",
    "\n",
    "            \"\"\"\n",
    "            return_dict = {\"data\": torch.from_numpy(batch.pop(\"data\")).to(\n",
    "                input_device).to(torch.float)}\n",
    "\n",
    "            for key, vals in batch.items():\n",
    "                return_dict[key] = torch.from_numpy(vals).to(output_device).to(\n",
    "                    torch.float)\n",
    "\n",
    "            return return_dict\n",
    "        \n",
    "```\n",
    "\n",
    "## Saving and loading\n",
    "Now that we have the ability to implement `delira`-suitable TorchScript models, we want to store them on disk and load them again, so that we don't have to retrain them every time we want to use them. These I/O functions are usually located in `delira.io`. \n",
    "\n",
    "### Saving\n",
    "Our saving function utilizes multiple functions: `torch.jit.save` to simply save the model (including it's graph) and the `save_checkpoint_torch` function implemented for the `PyTorch` backend to store the trainer state, since `TorchScript` allows us to use plain `PyTorch` optimizers.\n",
    "\n",
    "The implementation of the function looks like this:\n",
    "\n",
    "```python\n",
    "\n",
    "    def save_checkpoint_torchscript(file: str, model=None, optimizers={},\n",
    "                                    epoch=None, **kwargs):\n",
    "        \"\"\"\n",
    "        Save current checkpoint to two different files:\n",
    "            1.) ``file + \"_model.ptj\"``: Will include the state of the model\n",
    "                (including the graph; this is the opposite to\n",
    "                :func:`save_checkpoint`)\n",
    "            2.) ``file + \"_trainer_state.pt\"``: Will include the states of all\n",
    "                optimizers and the current epoch (if given)\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        file : str\n",
    "            filepath the model should be saved to\n",
    "        model : AbstractPyTorchJITNetwork or None\n",
    "            the model which should be saved\n",
    "            if None: empty dict will be saved as state dict\n",
    "        optimizers : dict\n",
    "            dictionary containing all optimizers\n",
    "        epoch : int\n",
    "            current epoch (will also be pickled)\n",
    "\n",
    "        \"\"\"\n",
    "\n",
    "        # remove file extension if given\n",
    "        if any([file.endswith(ext) for ext in [\".pth\", \".pt\", \".ptj\"]]):\n",
    "            file = file.rsplit(\".\", 1)[0]\n",
    "\n",
    "        if isinstance(model, AbstractPyTorchJITNetwork):\n",
    "            torch.jit.save(model, file + \"_model.ptj\")\n",
    "\n",
    "        if optimizers or epoch is not None:\n",
    "            save_checkpoint_torch(file + \"_trainer_state.pt\", None,\n",
    "                            optimizers=optimizers, epoch=epoch, **kwargs)\n",
    "            \n",
    "```\n",
    "\n",
    "### Loading\n",
    "To load a model, which has been saved to disk by this function we have to revert each part of it. We do this by using `torch.jit.load` for the model (and the graph) and `load_checkpoint_torch` by the `PyTorch` backend.\n",
    "The actual implementation is given here:\n",
    "\n",
    "```python\n",
    "\n",
    "    def load_checkpoint_torchscript(file: str, **kwargs):\n",
    "        \"\"\"\n",
    "        Loads a saved checkpoint consisting of 2 files\n",
    "        (see :func:`save_checkpoint_jit` for details)\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        file : str\n",
    "            filepath to a file containing a saved model\n",
    "        **kwargs:\n",
    "            Additional keyword arguments (passed to torch.load)\n",
    "            Especially \"map_location\" is important to change the device the\n",
    "            state_dict should be loaded to\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        OrderedDict\n",
    "            checkpoint state_dict\n",
    "\n",
    "        \"\"\"\n",
    "        # remove file extensions\n",
    "        if any([file.endswith(ext) for ext in [\".pth\", \".pt\", \".ptj\"]]):\n",
    "            file = file.rsplit(\".\", 1)[0]\n",
    "\n",
    "        # load model\n",
    "        if os.path.isfile(file + \".ptj\"):\n",
    "            model_file = file\n",
    "        elif os.path.isfile(file + \"_model.ptj\"):\n",
    "            model_file = file + \"_model.ptj\"\n",
    "        else:\n",
    "            raise ValueError(\"No Model File found for %s\" % file)\n",
    "\n",
    "        # load trainer state (if possible)\n",
    "        trainer_file = model_file.replace(\"_model.ptj\", \"_trainer_state.pt\")\n",
    "        if os.path.isfile(trainer_file):\n",
    "            trainer_state = load_checkpoint_torch(trainer_file, **kwargs)\n",
    "\n",
    "        else:\n",
    "            trainer_state = {\"optimizer\": {},\n",
    "                             \"epoch\": None}\n",
    "\n",
    "        trainer_state.update({\"model\": torch.jit.load(model_file)})\n",
    "\n",
    "        return trainer_state\n",
    "    \n",
    "```\n",
    "\n",
    "\n",
    "## A Trainer to train\n",
    "Now, that we can define and save/load our models, we want to train them. Luckily `delira` has already implemented a very modular backend-agnostic trainer (the `BaseNetworkTrainer`) and build upon this a `PyTorchNetworkTrainer`. Since the training process in PyTorch and TorchScript is nearly the same, we can just extend the `PyTorchNetworkTrainer`. Usually one would have to extend the `BaseNetworkTrainer` to provide some backend specific functions (like necessary initializations, optimizer setup, seeding etc.). To see how this is done, you could either have a look at the `PyTorchNetworkTrainer` or the `TfNetworkTrainer` for tensorflow, which are both following this principle. Usually the only stuff to completely change is the loading/saving behavior and the `_setup` function, which defines the backend-specific initialization. Some other functions may have to be extended (by implementing the extension and calling the parent-classes function).\n",
    "\n",
    "### Things to change:\n",
    "\n",
    "By Subclassing the `PyTorchNetworkTrainer` we have to change the following things:\n",
    "\n",
    "* The trainer's default arguments\n",
    "\n",
    "* The behavior for trying to resume a previous training\n",
    "\n",
    "* The saving, loading and updating behavior\n",
    "\n",
    "We will access this one by one:\n",
    "\n",
    "#### The Default Arguments\n",
    "\n",
    "We want to use `AbstractTorchScriptNetwork`s instead of `AbstractPyTorchNetwork`s here and we have to change the behavior if passing multiple GPUs, because currently Multi-GPU training is not supported by `TorchScript`.\n",
    "\n",
    "To do this: we implement the functions `__init__`, apply our changes and forward these changes to the call of the base-classes `__init__` like this (omitted docstrings for the sake of shortness):\n",
    "\n",
    "```python\n",
    "\n",
    "class TorchScriptNetworkTrainer(PyTorchNetworkTrainer):\n",
    "        def __init__(self,\n",
    "                     network: AbstractTorchScriptNetwork,\n",
    "                     save_path: str,\n",
    "                     key_mapping,\n",
    "                     losses=None,\n",
    "                     optimizer_cls=None,\n",
    "                     optimizer_params={},\n",
    "                     train_metrics={},\n",
    "                     val_metrics={},\n",
    "                     lr_scheduler_cls=None,\n",
    "                     lr_scheduler_params={},\n",
    "                     gpu_ids=[],\n",
    "                     save_freq=1,\n",
    "                     optim_fn=create_optims_default,\n",
    "                     logging_type=\"tensorboardx\",\n",
    "                     logging_kwargs={},\n",
    "                     fold=0,\n",
    "                     callbacks=[],\n",
    "                     start_epoch=1,\n",
    "                     metric_keys=None,\n",
    "                     convert_batch_to_npy_fn=convert_torch_tensor_to_npy,\n",
    "                     criterions=None,\n",
    "                     val_freq=1,\n",
    "                     **kwargs):\n",
    "            \n",
    "            if len(gpu_ids) > 1:\n",
    "                # only use first GPU due to\n",
    "                # https://github.com/pytorch/pytorch/issues/15421\n",
    "                gpu_ids = [gpu_ids[0]]\n",
    "                logging.warning(\"Multiple GPUs specified. Torch JIT currently \"\n",
    "                                \"supports only single-GPU training. \"\n",
    "                                \"Switching to use only the first GPU for now...\")\n",
    "\n",
    "            super().__init__(network=network, save_path=save_path,\n",
    "                             key_mapping=key_mapping, losses=losses,\n",
    "                             optimizer_cls=optimizer_cls,\n",
    "                             optimizer_params=optimizer_params,\n",
    "                             train_metrics=train_metrics,\n",
    "                             val_metrics=val_metrics,\n",
    "                             lr_scheduler_cls=lr_scheduler_cls,\n",
    "                             lr_scheduler_params=lr_scheduler_params,\n",
    "                             gpu_ids=gpu_ids, save_freq=save_freq,\n",
    "                             optim_fn=optim_fn, logging_type=logging_type,\n",
    "                             logging_kwargs=logging_kwargs, fold=fold,\n",
    "                             callbacks=callbacks,\n",
    "                             start_epoch=start_epoch, metric_keys=metric_keys,\n",
    "                             convert_batch_to_npy_fn=convert_batch_to_npy_fn,\n",
    "                             mixed_precision=False, mixed_precision_kwargs={},\n",
    "                             criterions=criterions, val_freq=val_freq, **kwargs\n",
    "                             )\n",
    "            \n",
    "```\n",
    "\n",
    "#### Resuming Training\n",
    "\n",
    "For resuming the training, we have to completely change the `try_resume_training` function and cannot reuse the parent's implementation of it. Thus, we don't call `super().try_resume_training` here, but completely reimplement it from scratch:\n",
    "\n",
    "```python\n",
    "\n",
    "    def try_resume_training(self):\n",
    "        \"\"\"\n",
    "        Load the latest state of a previous training if possible\n",
    "\n",
    "        \"\"\"\n",
    "        # Load latest epoch file if available\n",
    "        if os.path.isdir(self.save_path):\n",
    "            # check all files in directory starting with \"checkpoint\" and\n",
    "            # not ending with \"_best.pth\"\n",
    "            files = [x for x in os.listdir(self.save_path)\n",
    "                     if os.path.isfile(os.path.join(self.save_path, x))\n",
    "                     and x.startswith(\"checkpoint\")\n",
    "                     and not x.endswith(\"_best.ptj\")\n",
    "                     ]\n",
    "\n",
    "            # if list is not empty: load previous state\n",
    "            if files:\n",
    "\n",
    "                latest_epoch = max([\n",
    "                    int(x.rsplit(\"_\", 1)[-1].rsplit(\".\", 1)[0])\n",
    "                    for x in files])\n",
    "\n",
    "                latest_state_path = os.path.join(self.save_path,\n",
    "                                                 \"checkpoint_epoch_%d.ptj\"\n",
    "                                                 % latest_epoch)\n",
    "\n",
    "                # if pth file does not exist, load pt file instead\n",
    "                if not os.path.isfile(latest_state_path):\n",
    "                    latest_state_path = latest_state_path[:-1]\n",
    "\n",
    "                logger.info(\"Attempting to load state from previous \\\n",
    "                            training from %s\" % latest_state_path)\n",
    "                try:\n",
    "                    self.update_state(latest_state_path)\n",
    "                except KeyError:\n",
    "                    logger.warning(\"Previous State could not be loaded, \\\n",
    "                                    although it exists.Training will be \\\n",
    "                                    restarted\")\n",
    "\n",
    "```\n",
    "\n",
    "#### Saving and Loading\n",
    "Now we need to change the saving and loading behavior. As always we try to reuse as much code as possible to avoid code duplication.\n",
    "\n",
    "##### Saving\n",
    "To save the current training state, we simply call the `save_checkpoint_torchscript` function:\n",
    "\n",
    "```python\n",
    "\n",
    "    def save_state(self, file_name, epoch, **kwargs):\n",
    "        \"\"\"\n",
    "        saves the current state via\n",
    "        :func:`delira.io.torch.save_checkpoint_jit`\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        file_name : str\n",
    "            filename to save the state to\n",
    "        epoch : int\n",
    "            current epoch (will be saved for mapping back)\n",
    "        **kwargs :\n",
    "            keyword arguments\n",
    "\n",
    "        \"\"\"\n",
    "        if file_name.endswith(\".pt\") or file_name.endswith(\".pth\"):\n",
    "            file_name = file_name.rsplit(\".\", 1)[0]\n",
    "\n",
    "        save_checkpoint_torchscript(file_name, self.module, self.optimizers,\n",
    "                                    **kwargs)\n",
    "        \n",
    "```\n",
    "\n",
    "##### Loading\n",
    "\n",
    "To load the training state, we simply return the state loaded by `load_checkpoint_torchscript`.\n",
    "Since we don't use any arguments of the trainer itself here, the function is a `staticmethod`:\n",
    "\n",
    "```python\n",
    "\n",
    "    @staticmethod\n",
    "    def load_state(file_name, **kwargs):\n",
    "        \"\"\"\n",
    "        Loads the new state from file via\n",
    "        :func:`delira.io.torch.load_checkpoint:jit`\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        file_name : str\n",
    "            the file to load the state from\n",
    "        **kwargs : keyword arguments\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        dict\n",
    "            new state\n",
    "\n",
    "        \"\"\"\n",
    "        return load_checkpoint_torchscript(file_name, **kwargs)\n",
    "    \n",
    "```\n",
    "    \n",
    "##### Updating\n",
    "\n",
    "After we loaded the new state, we need to update the trainer's internal state by this new state.\n",
    "\n",
    "We do this by directly assigning the model here (since the graph was stored/loaded too) instead of only updating the state_dict and calling the parent-classes method afterwards:\n",
    "    \n",
    "```python\n",
    "\n",
    "    def _update_state(self, new_state):\n",
    "        \"\"\"\n",
    "        Update the state from a given new state\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        new_state : dict\n",
    "            new state to update internal state from\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        :class:`PyTorchNetworkJITTrainer`\n",
    "            the trainer with a modified state\n",
    "\n",
    "        \"\"\"\n",
    "        if \"model\" in new_state:\n",
    "            self.module = new_state.pop(\"model\").to(self.input_device)\n",
    "\n",
    "        return super()._update_state(new_state)\n",
    "\n",
    "```\n",
    " \n",
    "### A Whole Trainer\n",
    " \n",
    "After combining all the changes above, we finally get our new trainer as:\n",
    " \n",
    "```python\n",
    "\n",
    "    class TorchScriptNetworkTrainer(PyTorchNetworkTrainer):\n",
    "        def __init__(self,\n",
    "                     network: AbstractTorchScriptNetwork,\n",
    "                     save_path: str,\n",
    "                     key_mapping,\n",
    "                     losses=None,\n",
    "                     optimizer_cls=None,\n",
    "                     optimizer_params={},\n",
    "                     train_metrics={},\n",
    "                     val_metrics={},\n",
    "                     lr_scheduler_cls=None,\n",
    "                     lr_scheduler_params={},\n",
    "                     gpu_ids=[],\n",
    "                     save_freq=1,\n",
    "                     optim_fn=create_optims_default,\n",
    "                     logging_type=\"tensorboardx\",\n",
    "                     logging_kwargs={},\n",
    "                     fold=0,\n",
    "                     callbacks=[],\n",
    "                     start_epoch=1,\n",
    "                     metric_keys=None,\n",
    "                     convert_batch_to_npy_fn=convert_torch_tensor_to_npy,\n",
    "                     criterions=None,\n",
    "                     val_freq=1,\n",
    "                     **kwargs):\n",
    "            \"\"\"\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            network : :class:`AbstractPyTorchJITNetwork`\n",
    "                the network to train\n",
    "            save_path : str\n",
    "                path to save networks to\n",
    "            key_mapping : dict\n",
    "                a dictionary containing the mapping from the ``data_dict`` to\n",
    "                the actual model's inputs.\n",
    "                E.g. if a model accepts one input named 'x' and the data_dict\n",
    "                contains one entry named 'data' this argument would have to\n",
    "                be ``{'x': 'data'}``\n",
    "            losses : dict\n",
    "                dictionary containing the training losses\n",
    "            optimizer_cls : subclass of tf.train.Optimizer\n",
    "                optimizer class implementing the optimization algorithm of\n",
    "                choice\n",
    "            optimizer_params : dict\n",
    "                keyword arguments passed to optimizer during construction\n",
    "            train_metrics : dict, optional\n",
    "                metrics, which will be evaluated during train phase\n",
    "                (should work on framework's tensor types)\n",
    "            val_metrics : dict, optional\n",
    "                metrics, which will be evaluated during test phase\n",
    "                (should work on numpy arrays)\n",
    "            lr_scheduler_cls : Any\n",
    "                learning rate schedule class: must implement step() method\n",
    "            lr_scheduler_params : dict\n",
    "                keyword arguments passed to lr scheduler during construction\n",
    "            gpu_ids : list\n",
    "                list containing ids of GPUs to use; if empty: use cpu instead\n",
    "                Currently ``torch.jit`` only supports single GPU-Training,\n",
    "                thus only the first GPU will be used if multiple GPUs are passed\n",
    "            save_freq : int\n",
    "                integer specifying how often to save the current model's state.\n",
    "                State is saved every state_freq epochs\n",
    "            optim_fn : function\n",
    "                creates a dictionary containing all necessary optimizers\n",
    "            logging_type : str or callable\n",
    "                the type of logging. If string: it must be one of\n",
    "                [\"visdom\", \"tensorboardx\"]\n",
    "                If callable: it must be a logging handler class\n",
    "            logging_kwargs : dict\n",
    "                dictionary containing all logging keyword arguments\n",
    "            fold : int\n",
    "                current cross validation fold (0 per default)\n",
    "            callbacks : list\n",
    "                initial callbacks to register\n",
    "            start_epoch : int\n",
    "                epoch to start training at\n",
    "            metric_keys : dict\n",
    "                dict specifying which batch_dict entry to use for which metric as\n",
    "                target; default: None, which will result in key \"label\" for all\n",
    "                metrics\n",
    "            convert_batch_to_npy_fn : type, optional\n",
    "                function converting a batch-tensor to numpy, per default this is\n",
    "                a function, which detaches the tensor, moves it to cpu and the\n",
    "                calls ``.numpy()`` on it\n",
    "            mixed_precision : bool\n",
    "                whether to use mixed precision or not (False per default)\n",
    "            mixed_precision_kwargs : dict\n",
    "                additional keyword arguments for mixed precision\n",
    "            val_freq : int\n",
    "                validation frequency specifying how often to validate the trained\n",
    "                model (a value of 1 denotes validating every epoch,\n",
    "                a value of 2 denotes validating every second epoch etc.);\n",
    "                defaults to 1\n",
    "            **kwargs :\n",
    "                additional keyword arguments\n",
    "\n",
    "            \"\"\"\n",
    "\n",
    "            if len(gpu_ids) > 1:\n",
    "                # only use first GPU due to\n",
    "                # https://github.com/pytorch/pytorch/issues/15421\n",
    "                gpu_ids = [gpu_ids[0]]\n",
    "                logging.warning(\"Multiple GPUs specified. Torch JIT currently \"\n",
    "                                \"supports only single-GPU training. \"\n",
    "                                \"Switching to use only the first GPU for now...\")\n",
    "\n",
    "            super().__init__(network=network, save_path=save_path,\n",
    "                             key_mapping=key_mapping, losses=losses,\n",
    "                             optimizer_cls=optimizer_cls,\n",
    "                             optimizer_params=optimizer_params,\n",
    "                             train_metrics=train_metrics,\n",
    "                             val_metrics=val_metrics,\n",
    "                             lr_scheduler_cls=lr_scheduler_cls,\n",
    "                             lr_scheduler_params=lr_scheduler_params,\n",
    "                             gpu_ids=gpu_ids, save_freq=save_freq,\n",
    "                             optim_fn=optim_fn, logging_type=logging_type,\n",
    "                             logging_kwargs=logging_kwargs, fold=fold,\n",
    "                             callbacks=callbacks,\n",
    "                             start_epoch=start_epoch, metric_keys=metric_keys,\n",
    "                             convert_batch_to_npy_fn=convert_batch_to_npy_fn,\n",
    "                             mixed_precision=False, mixed_precision_kwargs={},\n",
    "                             criterions=criterions, val_freq=val_freq, **kwargs\n",
    "                             )\n",
    "\n",
    "        def try_resume_training(self):\n",
    "            \"\"\"\n",
    "            Load the latest state of a previous training if possible\n",
    "\n",
    "            \"\"\"\n",
    "            # Load latest epoch file if available\n",
    "            if os.path.isdir(self.save_path):\n",
    "                # check all files in directory starting with \"checkpoint\" and\n",
    "                # not ending with \"_best.pth\"\n",
    "                files = [x for x in os.listdir(self.save_path)\n",
    "                         if os.path.isfile(os.path.join(self.save_path, x))\n",
    "                         and x.startswith(\"checkpoint\")\n",
    "                         and not x.endswith(\"_best.ptj\")\n",
    "                         ]\n",
    "\n",
    "                # if list is not empty: load previous state\n",
    "                if files:\n",
    "\n",
    "                    latest_epoch = max([\n",
    "                        int(x.rsplit(\"_\", 1)[-1].rsplit(\".\", 1)[0])\n",
    "                        for x in files])\n",
    "\n",
    "                    latest_state_path = os.path.join(self.save_path,\n",
    "                                                     \"checkpoint_epoch_%d.ptj\"\n",
    "                                                     % latest_epoch)\n",
    "\n",
    "                    # if pth file does not exist, load pt file instead\n",
    "                    if not os.path.isfile(latest_state_path):\n",
    "                        latest_state_path = latest_state_path[:-1]\n",
    "\n",
    "                    logger.info(\"Attempting to load state from previous \\\n",
    "                                training from %s\" % latest_state_path)\n",
    "                    try:\n",
    "                        self.update_state(latest_state_path)\n",
    "                    except KeyError:\n",
    "                        logger.warning(\"Previous State could not be loaded, \\\n",
    "                                        although it exists.Training will be \\\n",
    "                                        restarted\")\n",
    "\n",
    "        def save_state(self, file_name, epoch, **kwargs):\n",
    "            \"\"\"\n",
    "            saves the current state via\n",
    "            :func:`delira.io.torch.save_checkpoint_jit`\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            file_name : str\n",
    "                filename to save the state to\n",
    "            epoch : int\n",
    "                current epoch (will be saved for mapping back)\n",
    "            **kwargs :\n",
    "                keyword arguments\n",
    "\n",
    "            \"\"\"\n",
    "            if file_name.endswith(\".pt\") or file_name.endswith(\".pth\"):\n",
    "                file_name = file_name.rsplit(\".\", 1)[0]\n",
    "\n",
    "            save_checkpoint_torchscript(file_name, self.module, self.optimizers,\n",
    "                                        **kwargs)\n",
    "\n",
    "        @staticmethod\n",
    "        def load_state(file_name, **kwargs):\n",
    "            \"\"\"\n",
    "            Loads the new state from file via\n",
    "            :func:`delira.io.torch.load_checkpoint:jit`\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            file_name : str\n",
    "                the file to load the state from\n",
    "            **kwargs : keyword arguments\n",
    "\n",
    "            Returns\n",
    "            -------\n",
    "            dict\n",
    "                new state\n",
    "\n",
    "            \"\"\"\n",
    "            return load_checkpoint_torchscript(file_name, **kwargs)\n",
    "\n",
    "        def _update_state(self, new_state):\n",
    "            \"\"\"\n",
    "            Update the state from a given new state\n",
    "\n",
    "            Parameters\n",
    "            ----------\n",
    "            new_state : dict\n",
    "                new state to update internal state from\n",
    "\n",
    "            Returns\n",
    "            -------\n",
    "            :class:`PyTorchNetworkJITTrainer`\n",
    "                the trainer with a modified state\n",
    "\n",
    "            \"\"\"\n",
    "            if \"model\" in new_state:\n",
    "                self.module = new_state.pop(\"model\").to(self.input_device)\n",
    "\n",
    "            return super()._update_state(new_state)\n",
    "        \n",
    "```\n",
    "\n",
    "## Wrapping it all in an Experiment\n",
    "To have access to methods like a K-Fold (and the not yet finished) hyperparameter tuning, we need to wrap the trainer in an Experiment. We will use the same approach as we did for implementing the trainer: Extending an already provided class.\n",
    "\n",
    "This time we extend the `PyTorchExperiment` which itself extends the `BaseExperiment` by some backend-specific defaults, types and seeds.\n",
    "\n",
    "Our whole class definition just changes the default arguments of the `PyTorchExperiment` and thus, we only have to implenent it's `__init__`:\n",
    "\n",
    "```python\n",
    "\n",
    "class TorchScriptExperiment(PyTorchExperiment):\n",
    "    def __init__(self,\n",
    "                 params: typing.Union[str, Parameters],\n",
    "                 model_cls: AbstractTorchScriptNetwork, # not AbstractPyTorchNetwork anymore\n",
    "                 n_epochs=None,\n",
    "                 name=None,\n",
    "                 save_path=None,\n",
    "                 key_mapping=None,\n",
    "                 val_score_key=None,\n",
    "                 optim_builder=create_optims_default_pytorch,\n",
    "                 checkpoint_freq=1,\n",
    "                 trainer_cls=TorchScriptNetworkTrainer, # not PyTorchNetworkTrainer anymore\n",
    "                 **kwargs):\n",
    "        \"\"\"\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        params : :class:`Parameters` or str\n",
    "            the training parameters, if string is passed,\n",
    "            it is treated as a path to a pickle file, where the\n",
    "            parameters are loaded from\n",
    "        model_cls : Subclass of :class:`AbstractTorchScriptNetwork`\n",
    "            the class implementing the model to train\n",
    "        n_epochs : int or None\n",
    "            the number of epochs to train, if None: can be specified later\n",
    "            during actual training\n",
    "        name : str or None\n",
    "            the Experiment's name\n",
    "        save_path : str or None\n",
    "            the path to save the results and checkpoints to.\n",
    "            if None: Current working directory will be used\n",
    "        key_mapping : dict\n",
    "            mapping between data_dict and model inputs (necessary for\n",
    "            prediction with :class:`Predictor`-API), if no keymapping is\n",
    "            given, a default key_mapping of {\"x\": \"data\"} will be used here\n",
    "        val_score_key : str or None\n",
    "            key defining which metric to use for validation (determining\n",
    "            best model and scheduling lr); if None: No validation-based\n",
    "            operations will be done (model might still get validated,\n",
    "            but validation metrics can only be logged and not used further)\n",
    "        optim_builder : function\n",
    "            Function returning a dict of backend-specific optimizers.\n",
    "            defaults to :func:`create_optims_default_pytorch`\n",
    "        checkpoint_freq : int\n",
    "            frequency of saving checkpoints (1 denotes saving every epoch,\n",
    "            2 denotes saving every second epoch etc.); default: 1\n",
    "        trainer_cls : subclass of :class:`TorchScriptNetworkTrainer`\n",
    "            the trainer class to use for training the model, defaults to\n",
    "            :class:`TorchScriptNetworkTrainer`\n",
    "        **kwargs :\n",
    "            additional keyword arguments\n",
    "\n",
    "        \"\"\"\n",
    "        super().__init__(params=params, model_cls=model_cls,\n",
    "                         n_epochs=n_epochs, name=name, save_path=save_path,\n",
    "                         key_mapping=key_mapping,\n",
    "                         val_score_key=val_score_key,\n",
    "                         optim_builder=optim_builder,\n",
    "                         checkpoint_freq=checkpoint_freq,\n",
    "                         trainer_cls=trainer_cls,\n",
    "                         **kwargs)\n",
    "        \n",
    "```\n",
    "\n",
    "## Testing it\n",
    "Now that we finished the implementation of the backend (which is the outermost wrapper; Congratulations!), we can just test it. We'll use a very simple network and test it with dummy data. We also only test the `run` and `test` functionality of our experiment, since everything else is just used for setting up the internal state or a composition of these two methods and already tested:\n",
    "Now, let's just define our dataset, instantiate it three times (for training, validation and testing) and wrap each of them into a `DataManager`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from delira.data_loading import AbstractDataset\n",
    "from delira.data_loading import DataManager\n",
    "\n",
    "\n",
    "class DummyDataset(AbstractDataset):\n",
    "    def __init__(self, length):\n",
    "        super().__init__(None, None)\n",
    "        self.length = length\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return {\"data\": np.random.rand(32),\n",
    "                \"label\": np.random.randint(0, 1, 1)}\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "\n",
    "    def get_sample_from_index(self, index):\n",
    "        return self.__getitem__(index)\n",
    "    \n",
    "dset_train = DummyDataset(500)\n",
    "dset_val = DummyDataset(50)\n",
    "dset_test = DummyDataset(10)\n",
    "\n",
    "# training, validation and testing with \n",
    "#a batchsize of 16, 1 loading thread and no transformations.\n",
    "dmgr_train = DataManager(dset_train, 16, 1, None)\n",
    "dmgr_val = DataManager(dset_val, 16, 1, None)\n",
    "dmgr_test = DataManager(dset_test, 16, 1, None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, that we have created three datasets, we need to define our small dummy network. We do this by subclassing `delira.models.AbstractTorchScriptNetwork` (which is the exactly implementation given above, be we need to use the internal one, because there are some typechecks against this one)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from delira.models import AbstractTorchScriptNetwork\n",
    "import torch\n",
    "\n",
    "\n",
    "class DummyNetworkTorchScript(AbstractTorchScriptNetwork):\n",
    "    __constants__ = [\"module\"]\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.module = self._build_model(32, 1)\n",
    "\n",
    "    @torch.jit.script_method\n",
    "    def forward(self, x):\n",
    "        return {\"pred\": self.module(x)}\n",
    "\n",
    "    @staticmethod\n",
    "    def prepare_batch(batch_dict, input_device, output_device):\n",
    "        return {\"data\": torch.from_numpy(batch_dict[\"data\"]\n",
    "                                         ).to(input_device,\n",
    "                                              torch.float),\n",
    "                \"label\": torch.from_numpy(batch_dict[\"label\"]\n",
    "                                          ).to(output_device,\n",
    "                                               torch.float)}\n",
    "\n",
    "    @staticmethod\n",
    "    def closure(model: AbstractTorchScriptNetwork, data_dict: dict,\n",
    "                optimizers: dict, losses={}, metrics={},\n",
    "                fold=0, **kwargs):\n",
    "        \"\"\"\n",
    "        closure method to do a single backpropagation step\n",
    "\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        model : \n",
    "            trainable model\n",
    "        data_dict : dict\n",
    "            dictionary containing the data\n",
    "        optimizers : dict\n",
    "            dictionary of optimizers to optimize model's parameters\n",
    "        losses : dict\n",
    "            dict holding the losses to calculate errors\n",
    "            (gradients from different losses will be accumulated)\n",
    "        metrics : dict\n",
    "            dict holding the metrics to calculate\n",
    "        fold : int\n",
    "            Current Fold in Crossvalidation (default: 0)\n",
    "        **kwargs:\n",
    "            additional keyword arguments\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        dict\n",
    "            Metric values (with same keys as input dict metrics)\n",
    "        dict\n",
    "            Loss values (with same keys as input dict losses)\n",
    "        list\n",
    "            Arbitrary number of predictions as torch.Tensor\n",
    "\n",
    "        Raises\n",
    "        ------\n",
    "        AssertionError\n",
    "            if optimizers or losses are empty or the optimizers are not\n",
    "            specified\n",
    "\n",
    "        \"\"\"\n",
    "\n",
    "        assert (optimizers and losses) or not optimizers, \\\n",
    "            \"Criterion dict cannot be emtpy, if optimizers are passed\"\n",
    "\n",
    "        loss_vals = {}\n",
    "        metric_vals = {}\n",
    "        total_loss = 0\n",
    "\n",
    "        # choose suitable context manager:\n",
    "        if optimizers:\n",
    "            context_man = torch.enable_grad\n",
    "\n",
    "        else:\n",
    "            context_man = torch.no_grad\n",
    "\n",
    "        with context_man():\n",
    "\n",
    "            inputs = data_dict.pop(\"data\")\n",
    "            preds = model(inputs)\n",
    "\n",
    "            if data_dict:\n",
    "\n",
    "                for key, crit_fn in losses.items():\n",
    "                    _loss_val = crit_fn(preds[\"pred\"], *data_dict.values())\n",
    "                    loss_vals[key] = _loss_val.item()\n",
    "                    total_loss += _loss_val\n",
    "\n",
    "                with torch.no_grad():\n",
    "                    for key, metric_fn in metrics.items():\n",
    "                        metric_vals[key] = metric_fn(\n",
    "                            preds[\"pred\"], *data_dict.values()).item()\n",
    "\n",
    "        if optimizers:\n",
    "            optimizers['default'].zero_grad()\n",
    "            # perform loss scaling via apex if half precision is enabled\n",
    "            with optimizers[\"default\"].scale_loss(total_loss) as scaled_loss:\n",
    "                scaled_loss.backward()\n",
    "            optimizers['default'].step()\n",
    "\n",
    "        else:\n",
    "\n",
    "            # add prefix \"val\" in validation mode\n",
    "            eval_loss_vals, eval_metrics_vals = {}, {}\n",
    "            for key in loss_vals.keys():\n",
    "                eval_loss_vals[\"val_\" + str(key)] = loss_vals[key]\n",
    "\n",
    "            for key in metric_vals:\n",
    "                eval_metrics_vals[\"val_\" + str(key)] = metric_vals[key]\n",
    "\n",
    "            loss_vals = eval_loss_vals\n",
    "            metric_vals = eval_metrics_vals\n",
    "\n",
    "        return metric_vals, loss_vals, {k: v.detach()\n",
    "                                        for k, v in preds.items()}\n",
    "\n",
    "    @staticmethod\n",
    "    def _build_model(in_channels, n_outputs):\n",
    "        return torch.nn.Sequential(\n",
    "            torch.nn.Linear(in_channels, 64),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(64, n_outputs)\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, that we defined our model, let's just test, if we really can forward some tensors through it. We will just use some random `torch.Tensors` (created by `torch.rand`). Since our model accepts 1d inputs of length 32, we need to pass 2d tensors to it (the additional dimension is the batch-dimension)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'single': tensor([[-0.1934]], grad_fn=<DifferentiableGraphBackward>),\n",
       " 'batched': tensor([[-0.0525],\n",
       "         [-0.0884],\n",
       "         [-0.1492],\n",
       "         [-0.0431]], grad_fn=<DifferentiableGraphBackward>)}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_tensor_single = torch.rand(1, 32) # use a single-sample batch (batchsize=1) here\n",
    "input_tensor_batched = torch.rand(4, 32) # use a batch with batchsize 4 here\n",
    "\n",
    "# create model instance\n",
    "model = DummyNetworkTorchScript()\n",
    "\n",
    "outputs = {\"single\": model(input_tensor_single)[\"pred\"], \"batched\": model(input_tensor_batched)[\"pred\"]}\n",
    "outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import mean_absolute_error\n",
    "from delira.training.callbacks import ReduceLROnPlateauCallbackPyTorch\n",
    "from delira.training import Parameters\n",
    "params = Parameters(fixed_params={\n",
    "                    \"model\": {},\n",
    "                    \"training\": {\n",
    "                        \"losses\": {\"CE\": torch.nn.BCEWithLogitsLoss()},\n",
    "                        \"optimizer_cls\": torch.optim.Adam,\n",
    "                        \"optimizer_params\": {\"lr\": 1e-3},\n",
    "                        \"num_epochs\": 2,\n",
    "                        \"val_metrics\": {\"mae\": mean_absolute_error},\n",
    "                        \"lr_sched_cls\": ReduceLROnPlateauCallbackPyTorch,\n",
    "                        \"lr_sched_params\": {\"mode\": \"min\"}\n",
    "                    }\n",
    "                }\n",
    "          )\n",
    "\n",
    "from delira.training import TorchScriptExperiment\n",
    "\n",
    "exp = TorchScriptExperiment(params, DummyNetworkTorchScript,\n",
    "                            key_mapping={\"x\": \"data\"},\n",
    "                            val_score_key=\"mae\",\n",
    "                            val_score_mode=\"min\")\n",
    "\n",
    "trained_model = exp.run(dmgr_train, dmgr_val)\n",
    "exp.test(trained_model, dmgr_test, params.nested_get(\"val_metrics\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations. You have implemented your first fully-workable `delira`-Backend. Wasn't that hard, was it?\n",
    "\n",
    "Before you start implementing backends for all the other frameworks out there, let me just give you some advices:\n",
    "\n",
    "* You should test everything you implement or extend\n",
    "\n",
    "* Make sure, to keep your backend-specification in mind\n",
    "\n",
    "* Always follow the API of already existing backends. If this is not possible: test this extensively\n",
    "\n",
    "* If you extend another backend (like we did here; we extended the `PyTorch`-backend for `TorchScript`), make sure, that the \"base-backend\" is always installed (best if they can only be installed together)\n",
    "\n",
    "* If you have questions regarding the implementation, don't hestiate to contact us."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}