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

packages = \
['perceiver', 'perceiver.data', 'perceiver.model', 'perceiver.scripts']

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

install_requires = \
['einops>=0.4.0,<0.5.0',
 'fairscale>=0.4.0,<0.5.0',
 'jsonargparse[signatures]>=4.7.0,<4.8.0',
 'lightning-bolts>=0.5.0,<0.6.0',
 'pytorch-lightning>=1.6.0,<1.7.0',
 'tokenizers>=0.12.0,<0.13.0',
 'torch>=1.11.0,<1.12.0',
 'torchmetrics>=0.8.0,<0.9.0',
 'torchtext>=0.12.0,<0.13.0',
 'torchvision>=0.12.0,<0.13.0']

setup_kwargs = {
    'name': 'perceiver-io',
    'version': '0.2.0',
    'description': 'Perceiver IO',
    'long_description': '# Perceiver IO\n\nThis project is a PyTorch implementation of\n\n- [Perceiver: General Perception with Iterative Attention](https://arxiv.org/abs/2103.03206) and\n- [Perceiver IO: A General Architecture for Structured Inputs & Outputs](https://arxiv.org/abs/2107.14795)\n\nand supports training of Perceiver and Perceiver IO models with [Pytorch Lightning](https://www.pytorchlightning.ai/)\nat any scale.\n\n- Section [Architecture](#architecture) explains how the architectures of Perceiver IO and Perceiver are implemented by\n  this project.\n- Section [Model API](#model-api) gives an example how a Perceiver IO model can be created and configured with the\n  [PyTorch model API](#pytorch-model-api), the [PyTorch Lightning model API](#pytorch-lightning-model-api) and the\n  [command line interface](#pytorch-lightning-model-cli).\n- Section [Training examples](#training-examples) demonstrates how Perceiver IO models can be trained on some example\n  tasks.\n- Section [Inference examples](#inference-examples) links to notebooks that demonstrate how trained Perceiver IO\n  models can be used for prediction.\n\n## Installation\n\n### From sources\n\n```shell\nconda env create -f environment.yml\nconda activate perceiver-io\npoetry install\n```\n\n### Via pip\n\n```shell\npip install perceiver-io\n```\n\nWhen installing via `pip` make sure you have a CUDA toolkit installed as well (see also [environment.yml](environment.yml)).\n\n## Architecture\n\nThe following sections describe how the conceptual architectures of Perceiver IO and Perceiver can be mapped to the\nimplementation provided by this project.\n\n### Perceiver IO\n\n![architecture-perceiver-io](docs/architecture-perceiver-io.png)\n\nNames of components shown in the implementation architecture are class names in the [PyTorch model API](#pytorch-model-api)\n(see also [model.py](perceiver/model/model.py)). Task-specific input and output adapters are subclasses of `InputAdapter`\nand `OuptutAdapter`, respectively (see also [adapter.py](perceiver/model/adapter.py)). Array dimensions (`M`, `C`),\n(`N`, `D`), (`O`, `F`)  and (`O`, `E`) have the following names in code and/or on the command line:\n\n| Array dimension | Configuration parameter name                                                    |\n|-----------------|---------------------------------------------------------------------------------|\n| `M`             | Input-specific name (e.g. `max_seq_len` for text input, ...)                    |\n| `C`             | `num_input_channels` (property of `InputAdapter`)                               |\n| `N`             | `num_latents`                                                                   |\n| `D`             | `num_latent_channels`                                                           |\n| `O`             | Output-specific name (e.g. `num_output_queries` for classification output, ...) |\n| `E`             | Output-specific name (e.g. `num_classes` for classification output, ...)        |\n| `F`             | `num_output_query_channels` (property of `OutputAdapter`)                       |\n\nThe number of layers in a `SelfAttentionBlock` can be specified with `num_self_attention_layers_per_block` and the\nnumber of blocks with `num_self_attention_blocks` (`L` in the conceptual architecture). Self-attention blocks share\ntheir weights.\n\n### Perceiver\n\nPerceiver IO does **not** use repeated encoder cross-attention as described the [Perceiver IO](https://arxiv.org/abs/2107.14795)\npaper:\n\n> We omit the repeated encoder cross-attends used in [Perceiver](https://arxiv.org/abs/2103.03206) as we found these to lead to relatively small performance\n> improvements but to significantly slow down training ...\n\nThis may be the case for the very large datasets used in the Perceiver IO paper but I found that repeated encoder\ncross-attention actually gives much better training results for smaller datasets. Therefore, the implementation\nprovided by this project supports repeated encoder cross-attention.\n\n![architecture-perceiver](docs/architecture-perceiver.png)\n\nThe number of repeated cross-attentions can be specified with `num_cross_attention_layers` (`P`) which must be less than\nor equal `num_self_attention_blocks` (`L`). Cross-attention layers 2 - `P` and self-attention blocks 2 - `L` always share\ntheir weights. Sharing the weights with the first cross-attention layer can be controlled with `first_cross_attention_layer_shared`,\nsharing the weights with the first self-attention block can be controlled with `first_self_attention_block_shared`. The\ndefault values of these hyperparameters are consistent with the Perceiver IO architecture (1 cross-attention layer, `L`\nself-attention blocks with weight sharing).\n\n## Model API\n\n### PyTorch model API\n\nThe PyTorch model API is based on generic encoder and decoder classes (`PerceiverEncoder` and `PerceiverDecoder`) and\ntask-specific input and output adapter classes. These are defined in [model.py](perceiver/model/model.py) and\n[adapter.py](perceiver/model/adapter.py), respectively. The following snippet shows how they can be used to create an\nImageNet classifier as specified in Appendix A of the [paper](https://arxiv.org/abs/2107.14795) (Perceiver IO, config A,\nwith 2D Fourier Features, 48.4M parameters):\n\n```python\nfrom perceiver.model import (\n    PerceiverIO,\n    PerceiverEncoder,\n    PerceiverDecoder,\n    ImageInputAdapter,\n    ClassificationOutputAdapter,\n)\n\n# Fourier-encodes pixel positions and flatten along spatial dimensions\ninput_adapter = ImageInputAdapter(\n    image_shape=(224, 224, 3),  # M = 224 * 224\n    num_frequency_bands=64,\n)\n\n# Projects generic Perceiver decoder output to specified number of classes\noutput_adapter = ClassificationOutputAdapter(\n    num_classes=1000,\n    num_output_query_channels=1024,  # F\n)  \n\n# Generic Perceiver encoder\nencoder = PerceiverEncoder(\n    input_adapter=input_adapter,\n    num_latents=512,  # N\n    num_latent_channels=1024,  # D\n    num_cross_attention_qk_channels=input_adapter.num_input_channels,  # C\n    num_cross_attention_heads=1,\n    num_self_attention_heads=8,\n    num_self_attention_layers_per_block=6,\n    num_self_attention_blocks=8,\n    dropout=0.0,\n)\n\n# Generic Perceiver decoder\ndecoder = PerceiverDecoder(\n    output_adapter=output_adapter,\n    num_latent_channels=1024,  # D\n    num_cross_attention_heads=1,\n    dropout=0.0,\n)\n\n# Perceiver IO image classifier\nmodel = PerceiverIO(encoder, decoder)\n```\n\n### PyTorch Lightning model API\n\nModels created with the [PyTorch model API](#pytorch-model-api) are wrapped in task-specific [LightningModule](https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html)s\n(e.g. `LitImageClassifier`) so that they can be trained with the PyTorch Lightning [Trainer](https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html).\nThey are defined in [lightning.py](perceiver/model/lightning.py). Part of this API are also task-specific configuration\nclasses defined in [config.py](perceiver/model/config.py).\n\nA task-specific encoder configuration class (e.g. `ImageEncoderConfig`) covers the configuration of the generic encoder\nand its task-specific input adapter. A task-specific decoder configuration object (e.g. `ClassificationDecoderConfig`)\ncovers the configuration of the generic decoder and its task-specific output adapter.\n\nThe same model as in the [previous section](#pytorch-model-api), wrapped in a `LitImageClassifier`, can be created\nwith:\n\n```python\nfrom perceiver.model.config import ImageEncoderConfig, ClassificationDecoderConfig\nfrom perceiver.model.lightning import LitImageClassifier\n\nencoder_cfg = ImageEncoderConfig(\n    image_shape=(224, 224, 3),\n    num_frequency_bands=64,\n    num_cross_attention_heads=1,\n    num_self_attention_heads=8,\n    num_self_attention_layers_per_block=6,\n    num_self_attention_blocks=8,\n    dropout=0.0,\n)\ndecoder_cfg = ClassificationDecoderConfig(\n    num_classes=1000,\n    num_output_query_channels=1024,\n    num_cross_attention_heads=1,\n    dropout=0.0,\n)\n\nlit_model = LitImageClassifier(encoder_cfg, decoder_cfg, num_latents=512, num_latent_channels=1024)\n\n# Wrapped PyTorch model\nmodel = lit_model.model\n```\n\n### PyTorch Lightning model CLI\n\nThe [PyTorch Lightning model API](#pytorch-lightning-model-api) is primarily designed for command-line binding via\nthe [Lightning CLI](https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_cli.html). For example, when\nimplementing a command line interface for `LitImageClassifier` with `LightningCLI` in a file named `classifier.py`\n\n```python\n# File classifier.py\n\nfrom pytorch_lightning.utilities.cli import LightningCLI\nfrom perceiver.model.lightning import LitImageClassifier\n\nif __name__ == "__main__":\n    LightningCLI(model_class=LitImageClassifier)\n```\n\nthe same classifier [as before](#pytorch-lightning-model-api) can be created with the following command line options:\n\n```shell\npython classifier.py fit \\\n  --model.num_latents=512 \\\n  --model.num_latent_channels=1024 \\\n  --model.encoder.image_shape=[224,224,3] \\\n  --model.encoder.num_frequency_bands=64 \\\n  --model.encoder.num_cross_attention_heads=1 \\\n  --model.encoder.num_self_attention_heads=8 \\\n  --model.encoder.num_self_attention_layers_per_block=6 \\\n  --model.encoder.num_self_attention_blocks=8 \\\n  --model.encoder.dropout=0.0 \\\n  --model.decoder.num_classes=1000 \\\n  --model.decoder.num_output_query_channels=1024 \\\n  --model.decoder.num_cross_attention_heads=1 \\\n  --model.decoder.dropout=0.0 \\\n  ...\n```\n\nTask-specific training scripts can set default values so that command lines are usually much shorter (see [img_clf.py](perceiver/scripts/img_clf.py)\nfor an example of a training script and section [Image classification](#image-classification) for a usage example).\n\n## Training examples\n\nIn the following subsections, Perceiver IO models are trained on a rather small scale (and on small datasets). In\nparticular, hyperparameters are set such that parallel training on two NVIDIA GTX 1080 GPUs (8 GB memory each) works\nquite well. I didn\'t really tune model architectures and other hyperparameters yet, so you\'ll probably get better\nresults with a bit of experimentation. Support for more datasets and tasks as well as instructions for training on\nlarger scale will come soon.\n\n### Masked language modeling\n\nPretrain a Perceiver IO model on masked language modeling (MLM) with text from the IMDB training set. The pretrained\nencoder is then used for training a [sentiment classification](#sentiment-classification) model.\n[Predictions of masked tokens](docs/tensorboard.md) are logged to Tensorboard.\n\n```shell\npython -m perceiver.scripts.mlm fit \\\n  --model.num_latents=64 \\\n  --model.num_latent_channels=64 \\\n  --model.encoder.num_input_channels=64 \\\n  --model.encoder.num_cross_attention_layers=3 \\\n  --model.encoder.num_self_attention_layers_per_block=6 \\\n  --model.encoder.num_self_attention_blocks=3 \\\n  --model.encoder.dropout=0.0 \\\n  --model.decoder.num_output_query_channels=64 \\\n  --model.decoder.dropout=0.0 \\\n  --data=ImdbDataModule \\\n  --data.max_seq_len=512 \\\n  --data.batch_size=64 \\\n  --optimizer.lr=3e-3 \\\n  --optimizer.weight_decay=0.0 \\\n  --lr_scheduler.pct_start=0.1 \\\n  --trainer.accelerator=gpu \\\n  --trainer.devices=-1 \\\n  --trainer.max_steps=50000 \\\n  --trainer.check_val_every_n_epoch=5 \\\n  --trainer.logger=TensorBoardLogger \\\n  --trainer.logger.save_dir=logs \\\n  --trainer.logger.name=mlm\n```\n\nFor saving GPU memory and scaling model training, [activation checkpointing](docs/checkpointing.md) can be enabled with\n`--model.activation_checkpointing=true` (disabled by default).\n\n### Sentiment classification\n\nTrain a classification decoder using a frozen encoder from [masked language modeling](#masked-language-modeling).\nIf you ran MLM yourself you\'ll need to modify the `--model.mlm_ckpt` argument accordingly, otherwise download\ncheckpoints from [here](https://martin-krasser.com/perceiver/logs-update-2.zip) and extract them in the root directory of\nthis project.\n\n```shell\npython -m perceiver.scripts.seq_clf fit \\\n  --model.mlm_ckpt=\'logs/mlm/version_0/checkpoints/epoch=254-val_loss=4.527.ckpt\' \\\n  --model.num_latents=64 \\\n  --model.num_latent_channels=64 \\\n  --model.encoder.num_input_channels=64 \\\n  --model.encoder.num_cross_attention_layers=3 \\\n  --model.encoder.num_self_attention_layers_per_block=6 \\\n  --model.encoder.num_self_attention_blocks=3 \\\n  --model.encoder.dropout=0.0 \\\n  --model.encoder.freeze=true \\\n  --model.decoder.num_output_query_channels=64 \\\n  --model.decoder.dropout=0.0 \\\n  --data=ImdbDataModule \\\n  --data.max_seq_len=512 \\\n  --data.batch_size=128 \\\n  --optimizer=AdamW \\\n  --optimizer.lr=1e-3 \\\n  --optimizer.weight_decay=0.01 \\\n  --trainer.accelerator=gpu \\\n  --trainer.devices=-1 \\\n  --trainer.max_epochs=30 \\\n  --trainer.logger=TensorBoardLogger \\\n  --trainer.logger.save_dir=logs \\\n  --trainer.logger.name=seq_clf\n```\n\nUnfreeze the encoder and jointly fine-tune it together with the decoder that has been trained in the previous step.\nIf you ran the previous step yourself you\'ll need to modify the `--model.clf_ckpt` argument accordingly, otherwise\ndownload checkpoints from [here](https://martin-krasser.com/perceiver/logs-update-2.zip).\n\n```shell\npython -m perceiver.scripts.seq_clf fit \\\n  --model.clf_ckpt=\'logs/seq_clf/version_0/checkpoints/epoch=009-val_loss=0.343.ckpt\' \\\n  --model.num_latents=64 \\\n  --model.num_latent_channels=64 \\\n  --model.encoder.num_input_channels=64 \\\n  --model.encoder.num_cross_attention_layers=3 \\\n  --model.encoder.num_self_attention_layers_per_block=6 \\\n  --model.encoder.num_self_attention_blocks=3 \\\n  --model.encoder.dropout=0.1 \\\n  --model.decoder.num_output_query_channels=64 \\\n  --model.decoder.dropout=0.1 \\\n  --data=ImdbDataModule \\\n  --data.max_seq_len=512 \\\n  --data.batch_size=128 \\\n  --optimizer=AdamW \\\n  --optimizer.lr=1e-4 \\\n  --optimizer.weight_decay=0.01 \\\n  --trainer.accelerator=gpu \\\n  --trainer.devices=-1 \\\n  --trainer.max_epochs=40 \\\n  --trainer.logger=TensorBoardLogger \\\n  --trainer.logger.save_dir=logs \\\n  --trainer.logger.name=seq_clf\n```\n\n### Image classification\n\nClassify MNIST images.\n\n```shell\npython -m perceiver.scripts.img_clf fit \\\n  --model.num_latents=32 \\\n  --model.num_latent_channels=128 \\\n  --model.encoder.num_self_attention_layers_per_block=3 \\\n  --model.encoder.num_self_attention_blocks=3 \\\n  --model.encoder.dropout=0.0 \\\n  --model.decoder.dropout=0.0 \\\n  --data=MnistDataModule \\\n  --data.batch_size=128 \\\n  --optimizer=AdamW \\\n  --optimizer.lr=1e-3 \\\n  --optimizer.weight_decay=0.01 \\\n  --trainer.accelerator=gpu \\\n  --trainer.devices=-1 \\\n  --trainer.max_epochs=20 \\\n  --trainer.logger=TensorBoardLogger \\\n  --trainer.logger.save_dir=logs \\\n  --trainer.logger.name=img_clf\n```\n\n## Inference examples\n\n- [Sentiment classification](https://colab.research.google.com/github/krasserm/perceiver-io/blob/main/notebooks/seq-clf.ipynb)\n- [Image classification](https://colab.research.google.com/github/krasserm/perceiver-io/blob/main/notebooks/img-clf.ipynb)\n\n## Development environment\n\nUpdate the project dependencies in the conda environment:\n\n```shell\ninvoke install\n```\n\nInstall the pre-commit hooks:\n\n```shell\ninvoke precommit-install\n```\n\nRun code quality checks:\n\n```shell\ninvoke cc\n```\n\nRun tests:\n\n```shell\ninvoke test\n```\n\nThe structure of this project is based on the [Python Project Template](https://github.com/cstub/python-project-template).\n\n## Citations\n\n```bibtex\n@misc{jaegle2021perceiver,\n    title   = {Perceiver: General Perception with Iterative Attention},\n    author  = {Andrew Jaegle and Felix Gimeno and Andrew Brock and Andrew Zisserman and Oriol Vinyals and Joao Carreira},\n    year    = {2021},\n    eprint  = {2103.03206},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CV}\n}\n```\n\n```bibtex\n@misc{jaegle2021perceiver,\n    title   = {Perceiver IO: A General Architecture for Structured Inputs & Outputs},\n    author  = {Andrew Jaegle and Sebastian Borgeaud and Jean-Baptiste Alayrac and Carl Doersch and Catalin Ionescu and David Ding and Skanda Koppula and Andrew Brock and Evan Shelhamer and Olivier Hénaff and Matthew M. Botvinick and Andrew Zisserman and Oriol Vinyals and João Carreira},\n    year    = {2021},\n    eprint  = {2107.14795},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.LG}\n}\n```\n',
    'author': 'Martin Krasser',
    'author_email': 'krasserm@googlemail.com',
    'maintainer': None,
    'maintainer_email': None,
    'url': 'https://github.com/krasserm/perceiver-io',
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'python_requires': '>=3.9.0,<3.10.0',
}


setup(**setup_kwargs)
