[SD-154] Create a PyTorch energy estimation CLI tool (#73)

* Add power plot for all pytorch models * Revert to pytorch raw data on DVC * Re-run pytorch training and experiments * Add ecoml CLI package * Fix failing tests * Update readme to include note about custom PyTorch models * Update Readme * Update introduction section in the README * Add tqdm progress when downloading models * Add case where there are no layers in the model * Handle error via rich error console * Update install instructions in README * Support ecoml in custom PyTorch workflow * chore: Formatting README * Fix path to model_summary in the README * Updated init function to load model --------- Co-authored-by: OCarrollM <[email protected]>
fuzzylabs · Feb 24, 2025 · 0c8d30b · 0c8d30b
1 parent c0ee0a8
commit 0c8d30b
Show file tree

Hide file tree

Showing 34 changed files with 8,684 additions and 42 deletions.
diff --git a/ecoml/.DS_Store b/ecoml/.DS_Store
diff --git a/ecoml/.gitignore b/ecoml/.gitignore
@@ -0,0 +1,182 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+ecoml_models
+
+# PyTorch models
+*.pth
+*.pt
+
+mlruns
diff --git a/ecoml/.python-version b/ecoml/.python-version
@@ -0,0 +1 @@
+3.12.3
diff --git a/ecoml/README.md b/ecoml/README.md
@@ -0,0 +1,56 @@
+# EcoML
+
+## Table of Contents
+
+- [Introduction](#introduction)
+- [Installation](#installation)
+- [Usage](#usage)
+
+## Introduction
+
+EcoML aims to provide accurate power consumption estimates for edge devices, helping developers optimize their applications for energy efficiency.
+
+## Installation
+
+To install the necessary dependencies, run the following command:
+
+```bash
+uv venv
+source .venv/bin/activate
+uv sync
+```
+
+> [!NOTE]
+> Coming Soon: We will also publish the package on PyPI for ease of use.
+
+## Usage
+
+To use EcoML for energy estimation of PyTorch models, follow these steps:
+
+1. Using PyTorch model summary
+
+    ```bash
+    ecoml predict --model sample_data/resnet18.json
+    ```
+
+    [Sample data](./sample_data/) folder contains model summary for 3 PyTorch models - Resnet18, Mobilenetv2 and VGG16.
+
+    To use a custom model for inference, you have to generate a model summary for the PyTorch model. Refer to the next section for how to use `ecoml` in your workflow.
+
+    `--verbose` flag can be passed to above command to get a detailed output.
+
+2. Using custom PyTorch model in your workflow
+
+    ```bash
+    from ecoml.model_summary.model_summary import get_summary
+
+    summary = get_summary(your_pt_model, model_input_shape, summary_file_path='summary/my_model.json')
+    ```
+
+    Here `your_pt_model` is a instance `nn.Module`, the trained PyTorch model.
+
+    Next, you can use the `predict` command to get the energy prediction using the path where model summary is saved.
+
+   ```bash
+    ecoml predict --model summary/my_model.json
+    ```
diff --git a/ecoml/pyproject.toml b/ecoml/pyproject.toml
@@ -0,0 +1,46 @@
+[project]
+name = "ecoml"
+version = "0.1.0"
+description = "Predict energy consumption of PyTorch CNN models on Jetson Orin"
+readme = "README.md"
+authors = [
+    { name = "fuzzylabs", email = "[email protected]" }
+]
+requires-python = ">=3.11"
+dependencies = [
+    "dagshub>=0.5.5",
+    "mlflow==2.18.0",
+    "pandas>=2.2.3",
+    "pydantic>=2.10.6",
+    "rich>=13.9.4",
+    "typer>=0.15.1",
+    "cloudpickle==3.1.0",
+    "psutil==6.1.0",
+    "scikit-learn==1.5.2",
+    "scipy==1.14.1",
+    "numpy==2.1.2",
+    "tqdm>=4.67.1",
+]
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+
+[tool.uv.sources]
+torch = [
+    { index = "pytorch-cpu", extra = "torch"},
+]
+
+[project.scripts]
+ecoml = "ecoml.ecoml:app"
+
+[project.optional-dependencies]
+torch = [
+    "torch>=2.6.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"