|
1 | 1 | # DeepQuant |
2 | 2 |
|
3 | | -A Python library for exporting Brevitas quantized neural networks. |
| 3 | +A library for true-quantization and optimization of neural networks. |
4 | 4 |
|
5 | 5 | ## Installation |
6 | 6 |
|
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: |
28 | 8 | ``` |
29 | | - |
30 | | -### Install the Package |
31 | | - |
32 | | -Clone the repository and install in development mode: |
33 | | - |
34 | | -```bash |
35 | | -cd DeepQuant |
36 | 9 | pip install -e . |
37 | 10 | ``` |
38 | 11 |
|
39 | 12 | ## Running Tests |
40 | 13 |
|
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`. |
109 | 15 |
|
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