Skip to content

Commit dcfa726

Browse files
committed
Update README
1 parent e1eadd0 commit dcfa726

File tree

1 file changed

+5
-198
lines changed

1 file changed

+5
-198
lines changed

README.md

Lines changed: 5 additions & 198 deletions
Original file line numberDiff line numberDiff line change
@@ -1,210 +1,17 @@
11
# DeepQuant
22

3-
A Python library for exporting Brevitas quantized neural networks.
3+
A library for true-quantization and optimization of neural networks.
44

55
## Installation
66

7-
### Requirements
8-
9-
- Python 3.11 or higher
10-
- PyTorch 2.1.2 or higher
11-
- Brevitas 0.11.0 or higher
12-
13-
### Setup Environment
14-
15-
First, create and activate a new conda environment:
16-
17-
```bash
18-
mamba create -n DeepQuant_env python=3.11
19-
mamba activate DeepQuant_env
20-
```
21-
22-
### Install Dependencies
23-
24-
Install PyTorch and its related packages:
25-
26-
```bash
27-
mamba install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 -c pytorch
7+
Start by creating a new env with `Python 3.11` or higher. Then clone the repo and install the library as an editable package with:
288
```
29-
30-
### Install the Package
31-
32-
Clone the repository and install in development mode:
33-
34-
```bash
35-
cd DeepQuant
369
pip install -e .
3710
```
3811

3912
## Running Tests
4013

41-
### Using Make (Recommended)
42-
43-
The project includes a Makefile with several testing commands:
44-
45-
```bash
46-
# Run all tests with verbose output
47-
make test
48-
49-
# Run only neural network test
50-
make test-nn
51-
52-
# Run only multi-head attention test
53-
make test-mha
54-
55-
# Run only CNN test
56-
make test-cnn
57-
58-
# Run only Resnet18 test
59-
make test-resnet
60-
61-
# Run a specific test file
62-
make test-single TEST=test_simple_nn.py
63-
64-
# Show all available make commands
65-
make help
66-
```
67-
68-
### Using pytest directly
69-
70-
You can also run tests using pytest commands:
71-
72-
```bash
73-
# Run all tests
74-
python -m pytest DeepQuant/tests -v -s
75-
76-
# Run a specific test file
77-
python -m pytest DeepQuant/tests/test_simple_nn.py -v -s
78-
```
79-
80-
## Project Structure
81-
82-
```
83-
DeepQuant/
84-
├── Makefile
85-
├── pyproject.toml
86-
├── conftest.py
87-
└── DeepQuant/
88-
├── custom_forwards/
89-
│ ├── activations.py
90-
│ ├── linear.py
91-
│ └── multiheadattention.py
92-
├── injects/
93-
│ ├── base.py
94-
│ ├── executor.py
95-
│ └── transformations.py
96-
├── tests/
97-
│ ├── test_simple_mha.py
98-
│ ├── test_simple_nn.py
99-
│ └── test_simple_cnn.py
100-
├── custom_tracer.py
101-
└── export_brevitas.py
102-
```
103-
104-
### Key Components
105-
106-
- **Makefile**: Provides automation commands for testing
107-
- **pyproject.toml**: Defines project metadata and dependencies for editable installation
108-
- **conftest.py**: Pytest configuration file that handles warning filters
14+
We provide comprehensive tests with pytest, to execute all tests, simply run `pytest`. We mark our tests in two categories, `SingleLayerTests` and `ModelTests`, to execute the tests of the specific category, you can run `pytest -m <category>`. For instance, to execute only the single layer tests, you can run `pytest -m SingleLayerTests`.
10915

110-
The source code is organized into several key modules:
111-
112-
- **custom_forwards/**: Contains the unrolled forward implementations for:
113-
114-
- Linear layers (QuantLinear, QuantConv2d)
115-
- Activation functions (QuantReLU, QuantSigmoid, etc.)
116-
- Multi-head attention (QuantMultiheadAttention)
117-
118-
- **injects/**: Contains the transformation infrastructure:
119-
120-
- Base transformation class and executor
121-
- Module-specific transformations
122-
- Validation and verification logic
123-
124-
- **tests/**: Example tests demonstrating the exporter usage:
125-
126-
- Simple neural network (linear + activations)
127-
- Multi-head attention model
128-
- Convolutional neural network
129-
- Resnet18
130-
131-
- **custom_tracer.py**: Implements a specialized `CustomBrevitasTracer` for FX tracing
132-
133-
- Handles Brevitas-specific module traversal
134-
- Ensures proper graph capture of quantization operations
135-
136-
- **export_brevitas.py**: Main API for end-to-end model export:
137-
- Orchestrates the transformation passes
138-
- Performs the final FX tracing
139-
- Validates model outputs through the process
140-
141-
## Usage
142-
143-
### Main Function: exportBrevitas
144-
145-
The main function of this library is `exportBrevitas`, which exports a Brevitas-based model to an FX GraphModule with unrolled quantization steps.
146-
147-
```python
148-
from DeepQuant.export_brevitas import exportBrevitas
149-
150-
# Initialize your Brevitas model
151-
model = YourBrevitasModel().eval()
152-
153-
# Create an input with the correct shape
154-
input = torch.randn(1, input_channels, height, width)
155-
156-
# Export the model (with debug information)
157-
fx_model = exportBrevitas(model, input, debug=True)
158-
```
159-
160-
Arguments:
161-
162-
- `model`: The Brevitas-based model to export
163-
- `example_input`: A representative input tensor for shape tracing
164-
- `debug`: If True, prints transformation progress (default: False)
165-
166-
When `debug=True`, you'll see the output showing the progress, for example:
167-
168-
```
169-
✓ MHA transformation successful - outputs match
170-
✓ Linear transformation successful - outputs match
171-
✓ Activation transformation successful - outputs match
172-
All transformations completed successfully!
173-
```
174-
175-
### Example Usage
176-
177-
A simple example script can be found in `example_usage.py` in the root directory of the project.
178-
179-
```python
180-
import torch
181-
import torch.nn as nn
182-
import brevitas.nn as qnn
183-
from brevitas.quant.scaled_int import Int8ActPerTensorFloat, Int32Bias
184-
from DeepQuant.export_brevitas import exportBrevitas
185-
186-
# Define a simple quantized model
187-
class SimpleQuantModel(nn.Module):
188-
def __init__(self):
189-
super().__init__()
190-
self.input_quant = qnn.QuantIdentity(return_quant_tensor=True)
191-
self.conv = qnn.QuantConv2d(
192-
in_channels=3,
193-
out_channels=16,
194-
kernel_size=3,
195-
bias=True,
196-
weight_bit_width=4,
197-
bias_quant=Int32Bias,
198-
output_quant=Int8ActPerTensorFloat,
199-
)
200-
201-
def forward(self, x):
202-
x = self.input_quant(x)
203-
x = self.conv(x)
204-
return x
205-
206-
# Export the model
207-
model = SimpleQuantModel().eval()
208-
dummy_input = torch.randn(1, 3, 32, 32)
209-
fx_model = exportBrevitas(model, dummy_input, debug=True)
210-
```
16+
## ⚠️ Disclaimer ⚠️
17+
This library is currently in **beta stage** and under active development. Interfaces and features are subject to change, and stability is not yet guaranteed. Use at your own risk, and feel free to report any issues or contribute to its improvement.

0 commit comments

Comments
 (0)