torchbench is now a library

.gitignore

@@ -19,4 +19,5 @@ old.json
 te.json
 logs/
 scripts/scribe.py
-.userbenchmark/
+.userbenchmark/
+torchbench.egg-info/

README.md

@@ -43,6 +43,30 @@ cd benchmark
 python install.py
 ```
 
+### Install torchbench as a library
+
+if you're interested in running torchbench as a library you can
+
+```bash
+python install.py
+pip install git+https://www.github.com:pytorch/benchmark.git
+```
+
+or 
+
+```bash
+python install.py
+pip install . # add -e for an editable installation
+```
+
+The above
+
+```python
+import torchbenchmark.models.densenet121
+model, example_inputs = torchbenchmark.models.densenet121.Model(test="eval", device="cuda", batch_size=1).get_module()
+model(*example_inputs)
+```
+
 ### Building From Source
 Note that when building PyTorch from source, torchvision and torchaudio must also be built from source to make sure the C APIs match.
 

setup.py

@@ -0,0 +1,16 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='torchbench',
+    version='0.1',
+    description='Benchmarking library for PyTorch',
+    author='PyTorch Team',
+    url='https://github.com/pytorch/benchmark',
+    packages=find_packages(include=['torchbenchmark*', 'userbenchmark*']),
+    classifiers=[
+        'Intended Audience :: Developers',
+        'Topic :: Software Development :: Build Tools',
+        'License :: OSI Approved :: BSD 3 License',
+        'Programming Language :: Python',
+    ],
+)

test_imports.py

@@ -0,0 +1,8 @@
+import torchbenchmark.models.densenet121
+model, example_inputs = torchbenchmark.models.densenet121.Model(test="eval", device="cuda", batch_size=1).get_module()
+model(*example_inputs)
+
+import userbenchmark.optim 
+import torchbenchmark.canary_models
+import torchbenchmark.models
+import torchbenchmark.score

torchbenchmark/__init__.py

@@ -17,6 +17,10 @@
 
 from components._impl.tasks import base as base_task
 from components._impl.workers import subprocess_worker
+from . import models
+from . import canary_models
+from . import e2e_models
+from . import util
 
 class ModelNotFoundError(RuntimeError):
     pass