# -*- coding: utf-8 -*-
from setuptools import setup

packages = \
['torchmtl']

package_data = \
{'': ['*']}

install_requires = \
['matplotlib>=3.3.2,<4.0.0',
 'networkx>=2.5,<3.0',
 'scipy>=1.5.3,<2.0.0',
 'torch>=1.6.0,<2.0.0']

setup_kwargs = {
    'name': 'torchmtl',
    'version': '0.1.8',
    'description': 'A lightweight module for Multi-Task Learning in pytorch',
    'long_description': '![torchMTL Logo](https://github.com/chrisby/torchMTL/blob/main/images/torchmtl_logo.png "torchMTL Logo")    \nA lightweight module for Multi-Task Learning in pytorch.\n\n`torchmtl` tries to help you composing modular multi-task architectures with minimal effort. All you need is a list of dictionaries in which you define your layers and how they build on each other. From this, `torchmtl` constructs a meta-computation graph which is executed in each forward pass of the created `MTLModel`. To combine outputs from multiple layers, simple [wrapper functions](https://github.com/chrisby/torchMTL/blob/main/torchmtl/wrapping_layers.py) are provided.\n\n### Installation\n`torchmtl` can be installed via `pip`:\n```\npip install torchmtl\n```\n\n### Quickstart (or find examples [here](https://github.com/chrisby/torchMTL/tree/main/examples))\nAssume you want to train a network on three tasks as shown below.  \n![example](https://github.com/chrisby/torchMTL/blob/main/images/example.png "example")  \n\nTo construct such an architecture with `torchmtl`, you simply have to define the following list\n\n```python\ntasks = [\n        {\n            \'name\': "Embed1",\n            \'layers\': Sequential(*[Linear(16, 32), Linear(32, 8)]),\n            # No anchor_layer means this layer receives input directly\n        },    \n        {\n            \'name\': "Embed2",\n            \'layers\': Sequential(*[Linear(16, 32), Linear(32, 8)]),\n            # No anchor_layer means this layer receives input directly\n        },\n        {\n            \'name\': "CatTask",\n            \'layers\': Concat(dim=1),\n            \'loss_weight\': 1.0,\n            \'anchor_layer\': [\'Embed1\', \'Embed2\']\n        },\n        {\n            \'name\': "Task1",\n            \'layers\': Sequential(*[Linear(8, 32), Linear(32, 1)]),\n            \'loss\': MSELoss(),\n            \'loss_weight\': 1.0,\n            \'anchor_layer\': \'Embed1\'            \n        },\n        {\n            \'name\': "Task2",\n            \'layers\': Sequential(*[Linear(8, 64), Linear(64, 1)]),\n            \'loss\': BCEWithLogitsLoss(),\n            \'loss_weight\': 1.0,\n            \'anchor_layer\': \'Embed2\'            \n        }, \n        {\n            \'name\': "FNN",\n            \'layers\': Sequential(*[Linear(16, 32), Linear(32, 32)]),\n            \'anchor_layer\': \'CatTask\'\n        },\n        {\n            \'name\': "Task3",\n            \'layers\': Sequential(*[Linear(32, 16), Linear(16, 1)]),\n            \'anchor_layer\': \'FNN\',\n            \'loss\': MSELoss(),\n            \'loss_weight\': \'auto\',\n            \'loss_init_val\': 1.0\n        }\n    ]\n```\n\nYou can build your final model with the following lines in which you specify from which layers you would like to receive the output.\n```python\nfrom torchmtl import MTLModel\nmodel = MTLModel(tasks, output_tasks=[\'Task1\', \'Task2\', \'Task3\'])\n```\n\nThis constructs a **meta-computation graph** which is executed in each forward pass of your `model`. You can verify whether the graph was properly built by plotting it using the `networkx` library:\n```python\nimport networkx as nx\npos = nx.planar_layout(model.g)\nnx.draw(model.g, pos, font_size=14, node_color="y", node_size=450, with_labels=True)\n```\n![graph example](https://github.com/chrisby/torchMTL/blob/main/images/torchmtl_graph.png "graph example")  \n\n#### The training loop\nYou can now enter the typical `pytorch` training loop and you will have access to everything you need to update your model:\n```python\nfor X, y in data_loader:\n    optimizer.zero_grad()\n\n    # Our model will return a list of predictions (from the layers specified in `output_tasks`),\n    # loss functions, and regularization parameters (as defined in the tasks variable)\n    y_hat, l_funcs, l_weights = model(X)\n    \n    loss = 0\n    # We can now iterate over the tasks and accumulate the losses\n    for i in range(len(y_hat)):\n        loss += l_weights[i] * l_funcs[i](y_hat[i], y[i])\n    \n    loss.backward()\n    optimizer.step()\n\n```\n\n### Details on the layer definition\nThere are 6 keys that can be specified (`name` and `layers` **must** always be present):  \n\n**`layers`**  \nBasically takes any `nn.Module` that you can think of. You can plug in a transformer or just a handful of fully connected layers.  \n\n**`anchor_layer`**  \nThis defines from which other layer this layer receives its input. Take care that the respective dimensions match.  \n\n**`loss`**  \nThe loss function you want to compute on the output of this layer (`l_funcs`). Can be set to `None` or omitted altogether when only access to the layer\'s output is needed.   \n\n**`loss_weight`**  \nThe scalar with which you want to regularize the respective loss (`l_weights`). If set to `\'auto\'`, a `nn.Parameter` is returned which will be updated through backpropagation. Can be set to `None` or omitted altogether when only access to the layer\'s output is needed.  \n\n**`loss_init_val`**  \nOnly needed if `loss_weight: \'auto\'`. The initialization value of the `loss_weight` parameter.\n\n### Wrapping functions\nNodes of the **meta-computation graph** don\'t have to be pytorch Modules. They can be *concatenation* functions or indexing functions that return a certain element of the input. If your `X` consists of two types of input data `X=[X_1, X_2]`, you can use the `SimpleSelect` layer to select the `X_1` by setting  \n```python\nfrom torchmtl.wrapping_layers import SimpleSelect\n{ ...,\n  \'layers\' = SimpleSelect(selection_axis=0),\n  ...\n}\n```\nIt should be trivial to write your own wrapping layers, but I try to provide useful ones with this library. If you have any layers in mind but no time to implement them, feel free to [open an issue](https://github.com/chrisby/torchMTL/issues).\n\nLogo credits and license: I reused and remixed (moved the dot and rotated the resulting logo a couple times) the pytorch logo from [here](https://github.com/pytorch/pytorch/blob/master/docs/source/_static/img/pytorch-logo-dark.png) (accessed through [wikimedia commons](https://commons.wikimedia.org/wiki/File:Pytorch_logo.png)) which can be used under the [Attribution-ShareAlike 4.0 International](https://creativecommons.org/licenses/by-sa/4.0/deed.en) license. Hence, this logo falls under the same license. \n',
    'author': 'Christian Bock',
    'author_email': 'christian.bock@bsse.ethz.ch',
    'maintainer': None,
    'maintainer_email': None,
    'url': None,
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'python_requires': '>=3.7,<4.0',
}


setup(**setup_kwargs)
