{"id":15600933,"url":"https://github.com/lucidrains/axial-attention","last_synced_at":"2025-04-04T21:09:39.747Z","repository":{"id":57413442,"uuid":"267681432","full_name":"lucidrains/axial-attention","owner":"lucidrains","description":"Implementation of Axial attention - attending to multi-dimensional data efficiently","archived":false,"fork":false,"pushed_at":"2021-08-26T01:14:39.000Z","size":39,"stargazers_count":375,"open_issues_count":5,"forks_count":33,"subscribers_count":11,"default_branch":"master","last_synced_at":"2025-03-28T20:11:26.815Z","etag":null,"topics":["artificial-intelligence","attention-mechanism","deep-learning","pytorch"],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2020-05-28T19:45:24.000Z","updated_at":"2025-03-28T18:33:24.000Z","dependencies_parsed_at":"2022-09-15T03:40:59.795Z","dependency_job_id":null,"html_url":"https://github.com/lucidrains/axial-attention","commit_stats":null,"previous_names":[],"tags_count":5,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Faxial-attention","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Faxial-attention/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Faxial-attention/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Faxial-attention/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/axial-attention/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247249530,"owners_count":20908212,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["artificial-intelligence","attention-mechanism","deep-learning","pytorch"],"created_at":"2024-10-03T02:09:40.264Z","updated_at":"2025-04-04T21:09:39.708Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"## Axial Attention\n\n[![PyPI version](https://badge.fury.io/py/axial-attention.svg)](https://badge.fury.io/py/axial-attention)\n\nImplementation of \u003ca href=\"https://arxiv.org/abs/1912.12180\"\u003eAxial attention\u003c/a\u003e in Pytorch. A simple but powerful technique to attend to multi-dimensional data efficiently. It has worked wonders for me and many other researchers.\n\nSimply add some positional encoding to your data and pass it into this handy class, specifying which dimension is considered the embedding, and how many axial dimensions to rotate through. All the permutating, reshaping, will be taken care of for you.\n\nThis paper was actually rejected on the basis of being too simple. And yet, it has since been used successfully in a number of applications, among those \u003ca href=\"https://ai.googleblog.com/2020/03/a-neural-weather-model-for-eight-hour.html\"\u003eweather prediction\u003c/a\u003e, \u003ca href=\"https://ai.googleblog.com/2020/08/axial-deeplab-long-range-modeling-in.html\"\u003e all-attention image segmentation\u003c/a\u003e. Just goes to show.\n\n### Install\n\n```bash\n$ pip install axial_attention\n```\n\n### Usage\n\nImage\n\n```python\nimport torch\nfrom axial_attention import AxialAttention\n\nimg = torch.randn(1, 3, 256, 256)\n\nattn = AxialAttention(\n    dim = 3,               # embedding dimension\n    dim_index = 1,         # where is the embedding dimension\n    dim_heads = 32,        # dimension of each head. defaults to dim // heads if not supplied\n    heads = 1,             # number of heads for multi-head attention\n    num_dimensions = 2,    # number of axial dimensions (images is 2, video is 3, or more)\n    sum_axial_out = True   # whether to sum the contributions of attention on each axis, or to run the input through them sequentially. defaults to true\n)\n\nattn(img) # (1, 3, 256, 256)\n```\n\nChannel-last image latents\n\n```python\nimport torch\nfrom axial_attention import AxialAttention\n\nimg = torch.randn(1, 20, 20, 512)\n\nattn = AxialAttention(\n    dim = 512,           # embedding dimension\n    dim_index = -1,      # where is the embedding dimension\n    heads = 8,           # number of heads for multi-head attention\n    num_dimensions = 2,  # number of axial dimensions (images is 2, video is 3, or more)\n)\n\nattn(img) # (1, 20, 20 ,512)\n```\n\nVideo\n\n```python\nimport torch\nfrom axial_attention import AxialAttention\n\nvideo = torch.randn(1, 5, 128, 256, 256)\n\nattn = AxialAttention(\n    dim = 128,           # embedding dimension\n    dim_index = 2,       # where is the embedding dimension\n    heads = 8,           # number of heads for multi-head attention\n    num_dimensions = 3,  # number of axial dimensions (images is 2, video is 3, or more)\n)\n\nattn(video) # (1, 5, 128, 256, 256)\n```\n\nImage Transformer, with reversible network\n\n```python\nimport torch\nfrom torch import nn\nfrom axial_attention import AxialImageTransformer\n\nconv1x1 = nn.Conv2d(3, 128, 1)\n\ntransformer = AxialImageTransformer(\n    dim = 128,\n    depth = 12,\n    reversible = True\n)\n\nimg = torch.randn(1, 3, 512, 512)\n\ntransformer(conv1x1(img)) # (1, 3, 512, 512)\n```\n\nWith axial positional embedding\n\n```python\nimport torch\nfrom axial_attention import AxialAttention, AxialPositionalEmbedding\n\nimg = torch.randn(1, 512, 20, 20)\n\nattn = AxialAttention(\n    dim = 512,\n    heads = 8,\n    dim_index = 1\n)\n\npos_emb = AxialPositionalEmbedding(\n    dim = 512,\n    shape = (20, 20)\n)\n\nimg = pos_emb(img)  # (1, 512, 20, 20)  - now positionally embedded\nimg = attn(img)     # (1, 512, 20, 20)\n```\n\n## Citation\n\n```bibtex\n@misc{ho2019axial,\n    title  = {Axial Attention in Multidimensional Transformers},\n    author = {Jonathan Ho and Nal Kalchbrenner and Dirk Weissenborn and Tim Salimans},\n    year   = {2019},\n    archivePrefix = {arXiv}\n}\n```\n\n```bibtex\n@misc{wang2020axialdeeplab,\n    title   = {Axial-DeepLab: Stand-Alone Axial-Attention for Panoptic Segmentation},\n    author  = {Huiyu Wang and Yukun Zhu and Bradley Green and Hartwig Adam and Alan Yuille and Liang-Chieh Chen},\n    year    = {2020},\n    eprint  = {2003.07853},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CV}\n}\n```\n\n```bibtex\n@inproceedings{huang2019ccnet,\n    title   = {Ccnet: Criss-cross attention for semantic segmentation},\n    author  = {Huang, Zilong and Wang, Xinggang and Huang, Lichao and Huang, Chang and Wei, Yunchao and Liu, Wenyu},\n    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision},\n    pages   = {603--612},\n    year    = {2019}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Faxial-attention","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Faxial-attention","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Faxial-attention/lists"}