Add initial project files

Author: federicobrancasi

Makefile

@@ -0,0 +1,54 @@
+# Copyright 2025 ETH Zurich and University of Bologna.
+# Licensed under the Apache License, Version 2.0, see LICENSE for details.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Federico Brancasi <fbrancasi@ethz.ch>
+
+.PHONY: test test-nn test-mha test-cnn
+
+# Test Directory
+TEST_DIR = src/DeepQuant/tests
+
+# Pytest flags
+PYTEST_FLAGS = -v -s
+
+# Target for running all tests
+test:
+	python -m pytest $(PYTEST_FLAGS) $(TEST_DIR)
+
+# Target for running simple neural network test
+test-nn:
+	python -m pytest $(PYTEST_FLAGS) $(TEST_DIR)/test_simple_nn.py
+
+# Target for running multi-head attention test
+test-mha:
+	python -m pytest $(PYTEST_FLAGS) $(TEST_DIR)/test_simple_mha.py
+
+# Target for running convolutional neural network test
+test-cnn:
+	python -m pytest $(PYTEST_FLAGS) $(TEST_DIR)/test_simple_cnn.py
+
+# Target for running resnet test
+test-resnet:
+	python -m pytest $(PYTEST_FLAGS) $(TEST_DIR)/test_resnet18.py
+
+# Target for running mnist test
+test-mnist:
+	python -m pytest $(PYTEST_FLAGS) $(TEST_DIR)/test_mnist.py
+
+# Target for running a specific test (usage: make test-single TEST=test_simple_nn.py)
+test-single:
+ifdef TEST
+	python -m pytest $(PYTEST_FLAGS) $(TEST_DIR)/$(TEST)
+else
+	@echo "Please specify a test file with TEST=filename.py"
+endif
+
+# Show help
+help:
+	@echo "Available targets:"
+	@echo "  make test      - Run all tests"
+	@echo "  make test-nn   - Run simple neural network tests"
+	@echo "  make test-mha  - Run multi-head attention tests"
+	@echo "  make test-cnn  - Run convolutional neural network tests"
+	@echo "  make test-single TEST=filename.py - Run a specific test file"

README.md

@@ -0,0 +1,211 @@
+# DeepQuant
+
+A Python library for exporting Brevitas quantized neural networks.
+
+## Installation
+
+### Requirements
+
+- Python 3.11 or higher
+- PyTorch 2.1.2 or higher
+- Brevitas 0.11.0 or higher
+
+### Setup Environment
+
+First, create and activate a new conda environment:
+
+```bash
+mamba create -n brevitas_env python=3.11
+mamba activate brevitas_env
+```
+
+### Install Dependencies
+
+Install PyTorch and its related packages:
+
+```bash
+mamba install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 -c pytorch
+```
+
+### Install the Package
+
+Clone the repository and install in development mode:
+
+```bash
+cd DeepQuant
+pip install -e .
+```
+
+## Running Tests
+
+### Using Make (Recommended)
+
+The project includes a Makefile with several testing commands:
+
+```bash
+# Run all tests with verbose output
+make test
+
+# Run only neural network test
+make test-nn
+
+# Run only multi-head attention test
+make test-mha
+
+# Run only CNN test
+make test-cnn
+
+# Run only Resnet18 test
+make test-resnet
+
+# Run a specific test file
+make test-single TEST=test_simple_nn.py
+
+# Show all available make commands
+make help
+```
+
+### Using pytest directly
+
+You can also run tests using pytest commands:
+
+```bash
+# Run all tests
+python -m pytest src/DeepQuant/tests -v -s
+
+# Run a specific test file
+python -m pytest src/DeepQuant/tests/test_simple_nn.py -v -s
+```
+
+## Project Structure
+
+```
+DeepQuant/
+├── Makefile
+├── pyproject.toml
+├── conftest.py
+└── src/
+    └── DeepQuant/
+        ├── custom_forwards/
+        │   ├── activations.py
+        │   ├── linear.py
+        │   └── multiheadattention.py
+        ├── injects/
+        │   ├── base.py
+        │   ├── executor.py
+        │   └── transformations.py
+        ├── tests/
+        │   ├── test_simple_mha.py
+        │   ├── test_simple_nn.py
+        │   └── test_simple_cnn.py
+        ├── custom_tracer.py
+        └── export_brevitas.py
+```
+
+### Key Components
+
+- **Makefile**: Provides automation commands for testing
+- **pyproject.toml**: Defines project metadata and dependencies for editable installation
+- **conftest.py**: Pytest configuration file that handles warning filters
+
+The source code is organized into several key modules:
+
+- **custom_forwards/**: Contains the unrolled forward implementations for:
+
+  - Linear layers (QuantLinear, QuantConv2d)
+  - Activation functions (QuantReLU, QuantSigmoid, etc.)
+  - Multi-head attention (QuantMultiheadAttention)
+
+- **injects/**: Contains the transformation infrastructure:
+
+  - Base transformation class and executor
+  - Module-specific transformations
+  - Validation and verification logic
+
+- **tests/**: Example tests demonstrating the exporter usage:
+
+  - Simple neural network (linear + activations)
+  - Multi-head attention model
+  - Convolutional neural network
+  - Resnet18
+
+- **custom_tracer.py**: Implements a specialized `CustomBrevitasTracer` for FX tracing
+
+  - Handles Brevitas-specific module traversal
+  - Ensures proper graph capture of quantization operations
+
+- **export_brevitas.py**: Main API for end-to-end model export:
+  - Orchestrates the transformation passes
+  - Performs the final FX tracing
+  - Validates model outputs through the process
+
+## Usage
+
+### Main Function: exportBrevitas
+
+The main function of this library is `exportBrevitas`, which exports a Brevitas-based model to an FX GraphModule with unrolled quantization steps.
+
+```python
+from DeepQuant.export_brevitas import exportBrevitas
+
+# Initialize your Brevitas model
+model = YourBrevitasModel().eval()
+
+# Create an input with the correct shape
+input = torch.randn(1, input_channels, height, width)
+
+# Export the model (with debug information)
+fx_model = exportBrevitas(model, input, debug=True)
+```
+
+Arguments:
+
+- `model`: The Brevitas-based model to export
+- `example_input`: A representative input tensor for shape tracing
+- `debug`: If True, prints transformation progress (default: False)
+
+When `debug=True`, you'll see the output showing the progress, for example:
+
+```
+✓ MHA transformation successful - outputs match
+✓ Linear transformation successful - outputs match
+✓ Activation transformation successful - outputs match
+All transformations completed successfully!
+```
+
+### Example Usage
+
+A simple example script can be found in `example_usage.py` in the root directory of the project.
+
+```python
+import torch
+import torch.nn as nn
+import brevitas.nn as qnn
+from brevitas.quant.scaled_int import Int8ActPerTensorFloat, Int32Bias
+from DeepQuant.export_brevitas import exportBrevitas
+
+# Define a simple quantized model
+class SimpleQuantModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.input_quant = qnn.QuantIdentity(return_quant_tensor=True)
+        self.conv = qnn.QuantConv2d(
+            in_channels=3,
+            out_channels=16,
+            kernel_size=3,
+            bias=True,
+            weight_bit_width=4,
+            bias_quant=Int32Bias,
+            output_quant=Int8ActPerTensorFloat,
+        )
+
+    def forward(self, x):
+        x = self.input_quant(x)
+        x = self.conv(x)
+        return x
+
+# Export the model
+model = SimpleQuantModel().eval()
+dummy_input = torch.randn(1, 3, 32, 32)
+fx_model = exportBrevitas(model, dummy_input, debug=True)
+```

__pycache__/conftest.cpython-311-pytest-8.3.4.pyc

@@ -0,0 +1,43 @@
+# Copyright 2025 ETH Zurich and University of Bologna.
+# Licensed under the Apache License, Version 2.0, see LICENSE for details.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Federico Brancasi <fbrancasi@ethz.ch>
+
+"""
+Pytest configuration file that suppresses specific warnings, including those
+related to torch.tensor constant registration in FX tracing.
+"""
+
+import warnings
+import pytest
+
+# Attempt to import TracerWarning from torch.fx.proxy;
+# if unavailable, skip filtering by category.
+try:
+    from torch.fx.proxy import TracerWarning
+
+    warnings.filterwarnings("ignore", category=TracerWarning)
+except ImportError:
+    pass
+
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+warnings.filterwarnings("ignore", category=UserWarning, message="Named tensors.*")
+warnings.filterwarnings(
+    "ignore", category=UserWarning, message=".*__torch_function__.*"
+)
+warnings.filterwarnings(
+    "ignore", category=UserWarning, message="Was not able to add assertion.*"
+)
+warnings.filterwarnings(
+    "ignore", category=UserWarning, message="'has_cuda' is deprecated.*"
+)
+warnings.filterwarnings(
+    "ignore", category=UserWarning, message="'has_cudnn' is deprecated.*"
+)
+warnings.filterwarnings(
+    "ignore", category=UserWarning, message="'has_mps' is deprecated.*"
+)
+warnings.filterwarnings(
+    "ignore", category=UserWarning, message="'has_mkldnn' is deprecated.*"
+)

__pycache__/conftest.cpython-311-pytest-8.3.5.pyc

@@ -0,0 +1,16 @@
+Metadata-Version: 2.2
+Name: custom_forwards
+Version: 0.1.0
+Summary: Custom PyTorch forwards library
+Requires-Python: >=3.11
+Requires-Dist: torch>=2.1.2
+Requires-Dist: torchvision>=0.16.2
+Requires-Dist: torchaudio>=2.1.2
+Requires-Dist: brevitas
+Requires-Dist: torchmetrics
+Provides-Extra: dev
+Requires-Dist: black; extra == "dev"
+Requires-Dist: isort; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: netron; extra == "dev"
+Requires-Dist: tabulate; extra == "dev"

conftest.py

@@ -0,0 +1,10 @@
+pyproject.toml
+setup.py
+custom_forwards.egg-info/PKG-INFO
+custom_forwards.egg-info/SOURCES.txt
+custom_forwards.egg-info/dependency_links.txt
+custom_forwards.egg-info/requires.txt
+custom_forwards.egg-info/top_level.txt
+src/__init__.py
+src/custom_tracer.py
+src/export_brevitas.py

custom_forwards.egg-info/PKG-INFO

@@ -0,0 +1 @@
+

custom_forwards.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+torch>=2.1.2
+torchvision>=0.16.2
+torchaudio>=2.1.2
+brevitas
+torchmetrics
+
+[dev]
+black
+isort
+pytest
+netron
+tabulate

custom_forwards.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+src