diff --git a/README.md b/README.md index 1db7943..efd7471 100644 --- a/README.md +++ b/README.md @@ -11,11 +11,11 @@ This implementation extends the [`timm`](https://github.com/huggingface/pytorch- ## Installation - - - +``` +pip install soft-moe +``` - +Or install the entire repo with: ``` git clone https://github.com/bwconrad/soft-moe @@ -29,7 +29,7 @@ pip install -r requirements.txt ```python import torch -from soft_moe_pytorch import SoftMoEVisionTransformer +from soft_moe import SoftMoEVisionTransformer net = SoftMoEVisionTransformer( num_experts=128, @@ -51,9 +51,9 @@ preds = net(img) Functions are also available to initialize default network configurations: ```python -from soft_moe_pytorch import (soft_moe_vit_base, soft_moe_vit_huge, - soft_moe_vit_large, soft_moe_vit_small, - soft_moe_vit_tiny) +from soft_moe import (soft_moe_vit_base, soft_moe_vit_huge, + soft_moe_vit_large, soft_moe_vit_small, + soft_moe_vit_tiny) net = soft_moe_vit_tiny() net = soft_moe_vit_small() @@ -95,7 +95,7 @@ The `SoftMoELayerWrapper` class can be used to make any network layer, that take import torch import torch.nn as nn -from soft_moe_pytorch import SoftMoELayerWrapper +from soft_moe import SoftMoELayerWrapper x = torch.rand(1, 16, 128) diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..c139f00 --- /dev/null +++ b/setup.py @@ -0,0 +1,36 @@ +from setuptools import find_packages, setup + +with open("README.md") as f: + long_description = f.read() + +setup( + name="soft_moe", + packages=find_packages(), + version="0.0.1", + license="Apache-2.0", + description="PyTorch implementation of 'From Sparse to Soft Mixtures of Experts'", + long_description=long_description, + long_description_content_type="text/markdown", + author="Ben Conrad", + author_email="benwconrad@proton.me", + url="https://github.com/bwconrad/soft-moe", + keywords=[ + "transformers", + "artificial intelligence", + "computer vision", + "deep learning", + ], + install_requires=[ + "timm >= 0.9.2", + "torch >= 2.0.1", + ], + classifiers=[ + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + ], + python_requires=">=3.10", +) diff --git a/soft_moe/__init__.py b/soft_moe/__init__.py new file mode 100644 index 0000000..eeade35 --- /dev/null +++ b/soft_moe/__init__.py @@ -0,0 +1,5 @@ +from soft_moe.soft_moe import SoftMoELayerWrapper +from soft_moe.vision_transformer import (SoftMoEVisionTransformer, + soft_moe_vit_base, soft_moe_vit_huge, + soft_moe_vit_large, + soft_moe_vit_small, soft_moe_vit_tiny) diff --git a/soft_moe_pytorch/soft_moe.py b/soft_moe/soft_moe.py similarity index 100% rename from soft_moe_pytorch/soft_moe.py rename to soft_moe/soft_moe.py diff --git a/soft_moe_pytorch/vision_transformer.py b/soft_moe/vision_transformer.py similarity index 100% rename from soft_moe_pytorch/vision_transformer.py rename to soft_moe/vision_transformer.py diff --git a/soft_moe_pytorch/__init__.py b/soft_moe_pytorch/__init__.py deleted file mode 100644 index 6ed06e5..0000000 --- a/soft_moe_pytorch/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .soft_moe import SoftMoELayerWrapper -from .vision_transformer import (SoftMoEVisionTransformer, soft_moe_vit_base, - soft_moe_vit_huge, soft_moe_vit_large, - soft_moe_vit_small, soft_moe_vit_tiny)