tracetorch-ml 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- tracetorch_ml-0.1.0/.github/workflows/ci.yml +31 -0
- tracetorch_ml-0.1.0/.github/workflows/publish.yml +17 -0
- tracetorch_ml-0.1.0/.gitignore +32 -0
- tracetorch_ml-0.1.0/.python-version +1 -0
- tracetorch_ml-0.1.0/LICENSE +21 -0
- tracetorch_ml-0.1.0/PKG-INFO +275 -0
- tracetorch_ml-0.1.0/README.md +247 -0
- tracetorch_ml-0.1.0/assets/header.jpg +0 -0
- tracetorch_ml-0.1.0/example.py +42 -0
- tracetorch_ml-0.1.0/pyproject.toml +61 -0
- tracetorch_ml-0.1.0/tests/__init__.py +0 -0
- tracetorch_ml-0.1.0/tests/test_tracetorch.py +606 -0
- tracetorch_ml-0.1.0/tracetorch/__init__.py +21 -0
- tracetorch_ml-0.1.0/tracetorch/analyzer.py +233 -0
- tracetorch_ml-0.1.0/tracetorch/cli.py +160 -0
- tracetorch_ml-0.1.0/tracetorch/collector.py +104 -0
- tracetorch_ml-0.1.0/tracetorch/hooks.py +177 -0
- tracetorch_ml-0.1.0/tracetorch/session.py +124 -0
- tracetorch_ml-0.1.0/tracetorch/storage.py +102 -0
- tracetorch_ml-0.1.0/tracetorch/utils.py +91 -0
- tracetorch_ml-0.1.0/uv.lock +939 -0
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
name: CI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [main]
|
|
6
|
+
pull_request:
|
|
7
|
+
branches: [main]
|
|
8
|
+
|
|
9
|
+
jobs:
|
|
10
|
+
lint:
|
|
11
|
+
runs-on: ubuntu-latest
|
|
12
|
+
steps:
|
|
13
|
+
- uses: actions/checkout@v4
|
|
14
|
+
- uses: astral-sh/setup-uv@v4
|
|
15
|
+
- run: uv venv && uv pip install ruff
|
|
16
|
+
- run: uv run ruff check .
|
|
17
|
+
- run: uv run ruff format --check .
|
|
18
|
+
|
|
19
|
+
test:
|
|
20
|
+
runs-on: ubuntu-latest
|
|
21
|
+
strategy:
|
|
22
|
+
matrix:
|
|
23
|
+
python-version: ["3.11", "3.12", "3.13"]
|
|
24
|
+
steps:
|
|
25
|
+
- uses: actions/checkout@v4
|
|
26
|
+
- uses: astral-sh/setup-uv@v4
|
|
27
|
+
with:
|
|
28
|
+
python-version: ${{ matrix.python-version }}
|
|
29
|
+
- run: uv venv && uv pip install -e ".[dev]"
|
|
30
|
+
- run: uv run pytest --tb=short -q
|
|
31
|
+
- run: uv run mypy tracetorch/
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
name: Publish to PyPI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
tags: ["v*"]
|
|
6
|
+
|
|
7
|
+
jobs:
|
|
8
|
+
publish:
|
|
9
|
+
runs-on: ubuntu-latest
|
|
10
|
+
permissions:
|
|
11
|
+
id-token: write
|
|
12
|
+
steps:
|
|
13
|
+
- uses: actions/checkout@v4
|
|
14
|
+
- uses: astral-sh/setup-uv@v4
|
|
15
|
+
- run: uv venv && uv pip install build
|
|
16
|
+
- run: uv run python -m build
|
|
17
|
+
- uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
# Python-generated files
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[oc]
|
|
4
|
+
build/
|
|
5
|
+
dist/
|
|
6
|
+
wheels/
|
|
7
|
+
*.egg-info
|
|
8
|
+
|
|
9
|
+
# Virtual environments
|
|
10
|
+
.venv
|
|
11
|
+
|
|
12
|
+
# Environment variables
|
|
13
|
+
.env
|
|
14
|
+
.env.*
|
|
15
|
+
|
|
16
|
+
# Testing
|
|
17
|
+
.pytest_cache/
|
|
18
|
+
.coverage
|
|
19
|
+
htmlcov/
|
|
20
|
+
|
|
21
|
+
# IDE
|
|
22
|
+
.vscode/
|
|
23
|
+
.idea/
|
|
24
|
+
*.swp
|
|
25
|
+
*.swo
|
|
26
|
+
|
|
27
|
+
# OS
|
|
28
|
+
.DS_Store
|
|
29
|
+
Thumbs.db
|
|
30
|
+
|
|
31
|
+
# Logs
|
|
32
|
+
*.log
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
3.13
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 rubai
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,275 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: tracetorch-ml
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Chrome DevTools for neural networks — interactive forensic debugging for PyTorch models
|
|
5
|
+
Project-URL: Homepage, https://github.com/phraakture/tracetorch
|
|
6
|
+
Project-URL: Repository, https://github.com/phraakture/tracetorch
|
|
7
|
+
Project-URL: Issues, https://github.com/phraakture/tracetorch/issues
|
|
8
|
+
License: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Intended Audience :: Developers
|
|
12
|
+
Classifier: Intended Audience :: Science/Research
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
18
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
19
|
+
Requires-Python: >=3.11
|
|
20
|
+
Requires-Dist: rich>=13.0
|
|
21
|
+
Requires-Dist: torch>=2.0
|
|
22
|
+
Provides-Extra: dev
|
|
23
|
+
Requires-Dist: mypy>=1.0; extra == 'dev'
|
|
24
|
+
Requires-Dist: pytest-cov>=4.0; extra == 'dev'
|
|
25
|
+
Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
26
|
+
Requires-Dist: ruff>=0.4; extra == 'dev'
|
|
27
|
+
Description-Content-Type: text/markdown
|
|
28
|
+
|
|
29
|
+
<p align="center">
|
|
30
|
+
<img src="assets/header.jpg?v=2" alt="TraceTorch" width="100%">
|
|
31
|
+
</p>
|
|
32
|
+
|
|
33
|
+
A forensic trace console for PyTorch models.
|
|
34
|
+
|
|
35
|
+
## Features
|
|
36
|
+
|
|
37
|
+
- **Zero-code tracing** -- wraps any `nn.Module` via context manager, no hook scattering
|
|
38
|
+
- **Per-layer capture** -- name, type, shape, dtype, device, params, latency
|
|
39
|
+
- **Activation statistics** -- mean, std, min, max, NaN count, Inf count per tensor
|
|
40
|
+
- **Backward gradient tracing** -- grad norm, mean, NaN detection, zero-gradient detection
|
|
41
|
+
- **Anomaly detection** -- 8 built-in checks: NaN, Inf, dead layers, exploding variance, variance spikes, zero gradients, NaN gradients, empty outputs
|
|
42
|
+
- **JSON export** -- structured, flat format designed for dashboards and CI pipelines
|
|
43
|
+
- **Rich CLI inspector** -- nested layer tree, anomaly warnings, model summary in the terminal
|
|
44
|
+
- **Accurate latency** -- pre-hook to post-hook timing, not just stat-computation overhead
|
|
45
|
+
|
|
46
|
+
<p>
|
|
47
|
+
<img src="https://img.shields.io/badge/python-3.11+-blue.svg" alt="Python">
|
|
48
|
+
<img src="https://img.shields.io/badge/pytorch-2.0+-ee4c2c.svg" alt="PyTorch">
|
|
49
|
+
<img src="https://img.shields.io/github/license/phraakture/tracetorch?v=1" alt="License">
|
|
50
|
+
</p>
|
|
51
|
+
|
|
52
|
+
## Install
|
|
53
|
+
|
|
54
|
+
```bash
|
|
55
|
+
pip install tracetorch
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
Requires Python 3.11+. Runtime dependencies: PyTorch 2.0+, Rich 13.0+.
|
|
59
|
+
|
|
60
|
+
## Quick start
|
|
61
|
+
|
|
62
|
+
```python
|
|
63
|
+
import torch
|
|
64
|
+
import torch.nn as nn
|
|
65
|
+
from tracetorch import TraceSession
|
|
66
|
+
|
|
67
|
+
model = nn.Sequential(
|
|
68
|
+
nn.Linear(784, 256),
|
|
69
|
+
nn.ReLU(),
|
|
70
|
+
nn.Linear(256, 10),
|
|
71
|
+
)
|
|
72
|
+
|
|
73
|
+
session = TraceSession(model)
|
|
74
|
+
|
|
75
|
+
with session:
|
|
76
|
+
x = torch.randn(32, 784)
|
|
77
|
+
y = model(x)
|
|
78
|
+
|
|
79
|
+
print(session.summary())
|
|
80
|
+
session.export("trace.json")
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
Inspect the saved trace:
|
|
84
|
+
|
|
85
|
+
```bash
|
|
86
|
+
tracetorch inspect trace.json
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
## What gets captured
|
|
90
|
+
|
|
91
|
+
For each child module reached through `model.named_modules()`:
|
|
92
|
+
|
|
93
|
+
- module name, full name, type, and nesting depth
|
|
94
|
+
- direct parameter count (`recurse=False`)
|
|
95
|
+
- input and output tensor shape, dtype, and device
|
|
96
|
+
- floating tensor statistics: mean, std, min, max, NaN count, Inf count
|
|
97
|
+
- forward latency measured from pre-hook to post-hook
|
|
98
|
+
- backward gradient norm, mean, NaN flag, zero-gradient flag
|
|
99
|
+
|
|
100
|
+
Root model is not traced; only its children are.
|
|
101
|
+
|
|
102
|
+
## Backward tracing
|
|
103
|
+
|
|
104
|
+
Run backward inside the session to capture gradient data:
|
|
105
|
+
|
|
106
|
+
```python
|
|
107
|
+
session = TraceSession(model)
|
|
108
|
+
|
|
109
|
+
with session:
|
|
110
|
+
x = torch.randn(4, 10, requires_grad=True)
|
|
111
|
+
out = model(x)
|
|
112
|
+
loss = out.sum()
|
|
113
|
+
loss.backward()
|
|
114
|
+
|
|
115
|
+
for layer in session.record.layers:
|
|
116
|
+
print(layer.full_name, layer.grad_norm, layer.has_zero_grad)
|
|
117
|
+
```
|
|
118
|
+
|
|
119
|
+
Gradient fields stay `None` / `False` if no backward pass runs.
|
|
120
|
+
|
|
121
|
+
## Anomaly checks
|
|
122
|
+
|
|
123
|
+
| Type | Trigger |
|
|
124
|
+
| --- | --- |
|
|
125
|
+
| `nan_activation` | Output tensor contains NaN |
|
|
126
|
+
| `inf_activation` | Output tensor contains Inf |
|
|
127
|
+
| `dead_layer` | Output mean and std both near zero |
|
|
128
|
+
| `exploding_variance` | Std > 100, or >10x spike from previous layer |
|
|
129
|
+
| `high_variance` | Std between 10 and 100 |
|
|
130
|
+
| `empty_output` | No tensor output captured |
|
|
131
|
+
| `zero_gradient` | Backward gradient is all zeros |
|
|
132
|
+
| `nan_gradient` | Backward gradient contains NaN |
|
|
133
|
+
|
|
134
|
+
## JSON format
|
|
135
|
+
|
|
136
|
+
```json
|
|
137
|
+
{
|
|
138
|
+
"metadata": {
|
|
139
|
+
"model_name": "Sequential",
|
|
140
|
+
"pytorch_version": "2.x",
|
|
141
|
+
"python_version": "3.11.x",
|
|
142
|
+
"cuda_available": false,
|
|
143
|
+
"active_device": "cpu",
|
|
144
|
+
"total_time_ms": 1.23
|
|
145
|
+
},
|
|
146
|
+
"layers": [
|
|
147
|
+
{
|
|
148
|
+
"name": "0",
|
|
149
|
+
"type": "Linear",
|
|
150
|
+
"full_name": "0",
|
|
151
|
+
"depth": 1,
|
|
152
|
+
"params": 200960,
|
|
153
|
+
"inputs": [
|
|
154
|
+
{
|
|
155
|
+
"shape": [32, 784],
|
|
156
|
+
"dtype": "torch.float32",
|
|
157
|
+
"device": "cpu",
|
|
158
|
+
"mean": 0.0,
|
|
159
|
+
"std": 1.0,
|
|
160
|
+
"min": -3.0,
|
|
161
|
+
"max": 3.0,
|
|
162
|
+
"nan_count": 0,
|
|
163
|
+
"inf_count": 0
|
|
164
|
+
}
|
|
165
|
+
],
|
|
166
|
+
"outputs": [
|
|
167
|
+
{
|
|
168
|
+
"shape": [32, 256],
|
|
169
|
+
"dtype": "torch.float32",
|
|
170
|
+
"device": "cpu",
|
|
171
|
+
"mean": 0.01,
|
|
172
|
+
"std": 0.58,
|
|
173
|
+
"min": -2.1,
|
|
174
|
+
"max": 2.4,
|
|
175
|
+
"nan_count": 0,
|
|
176
|
+
"inf_count": 0
|
|
177
|
+
}
|
|
178
|
+
],
|
|
179
|
+
"latency_ms": 0.25,
|
|
180
|
+
"grad_norm": null,
|
|
181
|
+
"grad_mean": null,
|
|
182
|
+
"grad_has_nan": false,
|
|
183
|
+
"has_zero_grad": false
|
|
184
|
+
}
|
|
185
|
+
],
|
|
186
|
+
"warnings": []
|
|
187
|
+
}
|
|
188
|
+
```
|
|
189
|
+
|
|
190
|
+
Load a saved trace in Python:
|
|
191
|
+
|
|
192
|
+
```python
|
|
193
|
+
from tracetorch.storage import load_json
|
|
194
|
+
|
|
195
|
+
record = load_json("trace.json")
|
|
196
|
+
print(record.metadata)
|
|
197
|
+
print(record.layers[0].outputs)
|
|
198
|
+
```
|
|
199
|
+
|
|
200
|
+
## API Reference
|
|
201
|
+
|
|
202
|
+
### TraceSession
|
|
203
|
+
|
|
204
|
+
```python
|
|
205
|
+
session = TraceSession(model, capture_stats=True, model_name=None)
|
|
206
|
+
```
|
|
207
|
+
|
|
208
|
+
Context manager for one traced execution window.
|
|
209
|
+
|
|
210
|
+
| Property / Method | Description |
|
|
211
|
+
| --- | --- |
|
|
212
|
+
| `record` | Collected `TraceRecord` |
|
|
213
|
+
| `anomalies` | `list[Anomaly]` populated after exit |
|
|
214
|
+
| `export(path)` | Write JSON, return `Path` |
|
|
215
|
+
| `summary()` | Compact text summary |
|
|
216
|
+
|
|
217
|
+
### Data structures
|
|
218
|
+
|
|
219
|
+
```python
|
|
220
|
+
from tracetorch import TraceRecord, LayerTrace, TensorStats, Anomaly
|
|
221
|
+
```
|
|
222
|
+
|
|
223
|
+
All dataclasses serialize via `.to_dict()`. JSON format is flat for dashboard consumption.
|
|
224
|
+
|
|
225
|
+
### CLI
|
|
226
|
+
|
|
227
|
+
```bash
|
|
228
|
+
tracetorch inspect trace.json
|
|
229
|
+
```
|
|
230
|
+
|
|
231
|
+
Prints model summary, nested layer tree with tensor shapes/params/latency/std, and anomaly warnings.
|
|
232
|
+
|
|
233
|
+
## Architecture
|
|
234
|
+
|
|
235
|
+
```
|
|
236
|
+
PyTorch Model
|
|
237
|
+
│
|
|
238
|
+
PyTorch Hooks
|
|
239
|
+
(forward_pre, forward, backward)
|
|
240
|
+
│
|
|
241
|
+
Trace Collector
|
|
242
|
+
capture shapes, stats, grads
|
|
243
|
+
│
|
|
244
|
+
Anomaly Detector
|
|
245
|
+
8 built-in checks
|
|
246
|
+
│
|
|
247
|
+
Storage Layer
|
|
248
|
+
JSON export / import
|
|
249
|
+
│
|
|
250
|
+
CLI Inspector
|
|
251
|
+
rich terminal output
|
|
252
|
+
```
|
|
253
|
+
|
|
254
|
+
## Development
|
|
255
|
+
|
|
256
|
+
```bash
|
|
257
|
+
git clone https://github.com/phraakture/tracetorch.git
|
|
258
|
+
cd tracetorch
|
|
259
|
+
pip install -e ".[dev]"
|
|
260
|
+
|
|
261
|
+
pytest # tests
|
|
262
|
+
ruff check tracetorch tests # lint
|
|
263
|
+
mypy tracetorch # types
|
|
264
|
+
```
|
|
265
|
+
|
|
266
|
+
Run the example:
|
|
267
|
+
|
|
268
|
+
```bash
|
|
269
|
+
python example.py
|
|
270
|
+
tracetorch inspect trace.json
|
|
271
|
+
```
|
|
272
|
+
|
|
273
|
+
## License
|
|
274
|
+
|
|
275
|
+
MIT
|
|
@@ -0,0 +1,247 @@
|
|
|
1
|
+
<p align="center">
|
|
2
|
+
<img src="assets/header.jpg?v=2" alt="TraceTorch" width="100%">
|
|
3
|
+
</p>
|
|
4
|
+
|
|
5
|
+
A forensic trace console for PyTorch models.
|
|
6
|
+
|
|
7
|
+
## Features
|
|
8
|
+
|
|
9
|
+
- **Zero-code tracing** -- wraps any `nn.Module` via context manager, no hook scattering
|
|
10
|
+
- **Per-layer capture** -- name, type, shape, dtype, device, params, latency
|
|
11
|
+
- **Activation statistics** -- mean, std, min, max, NaN count, Inf count per tensor
|
|
12
|
+
- **Backward gradient tracing** -- grad norm, mean, NaN detection, zero-gradient detection
|
|
13
|
+
- **Anomaly detection** -- 8 built-in checks: NaN, Inf, dead layers, exploding variance, variance spikes, zero gradients, NaN gradients, empty outputs
|
|
14
|
+
- **JSON export** -- structured, flat format designed for dashboards and CI pipelines
|
|
15
|
+
- **Rich CLI inspector** -- nested layer tree, anomaly warnings, model summary in the terminal
|
|
16
|
+
- **Accurate latency** -- pre-hook to post-hook timing, not just stat-computation overhead
|
|
17
|
+
|
|
18
|
+
<p>
|
|
19
|
+
<img src="https://img.shields.io/badge/python-3.11+-blue.svg" alt="Python">
|
|
20
|
+
<img src="https://img.shields.io/badge/pytorch-2.0+-ee4c2c.svg" alt="PyTorch">
|
|
21
|
+
<img src="https://img.shields.io/github/license/phraakture/tracetorch?v=1" alt="License">
|
|
22
|
+
</p>
|
|
23
|
+
|
|
24
|
+
## Install
|
|
25
|
+
|
|
26
|
+
```bash
|
|
27
|
+
pip install tracetorch
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
Requires Python 3.11+. Runtime dependencies: PyTorch 2.0+, Rich 13.0+.
|
|
31
|
+
|
|
32
|
+
## Quick start
|
|
33
|
+
|
|
34
|
+
```python
|
|
35
|
+
import torch
|
|
36
|
+
import torch.nn as nn
|
|
37
|
+
from tracetorch import TraceSession
|
|
38
|
+
|
|
39
|
+
model = nn.Sequential(
|
|
40
|
+
nn.Linear(784, 256),
|
|
41
|
+
nn.ReLU(),
|
|
42
|
+
nn.Linear(256, 10),
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
session = TraceSession(model)
|
|
46
|
+
|
|
47
|
+
with session:
|
|
48
|
+
x = torch.randn(32, 784)
|
|
49
|
+
y = model(x)
|
|
50
|
+
|
|
51
|
+
print(session.summary())
|
|
52
|
+
session.export("trace.json")
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
Inspect the saved trace:
|
|
56
|
+
|
|
57
|
+
```bash
|
|
58
|
+
tracetorch inspect trace.json
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
## What gets captured
|
|
62
|
+
|
|
63
|
+
For each child module reached through `model.named_modules()`:
|
|
64
|
+
|
|
65
|
+
- module name, full name, type, and nesting depth
|
|
66
|
+
- direct parameter count (`recurse=False`)
|
|
67
|
+
- input and output tensor shape, dtype, and device
|
|
68
|
+
- floating tensor statistics: mean, std, min, max, NaN count, Inf count
|
|
69
|
+
- forward latency measured from pre-hook to post-hook
|
|
70
|
+
- backward gradient norm, mean, NaN flag, zero-gradient flag
|
|
71
|
+
|
|
72
|
+
Root model is not traced; only its children are.
|
|
73
|
+
|
|
74
|
+
## Backward tracing
|
|
75
|
+
|
|
76
|
+
Run backward inside the session to capture gradient data:
|
|
77
|
+
|
|
78
|
+
```python
|
|
79
|
+
session = TraceSession(model)
|
|
80
|
+
|
|
81
|
+
with session:
|
|
82
|
+
x = torch.randn(4, 10, requires_grad=True)
|
|
83
|
+
out = model(x)
|
|
84
|
+
loss = out.sum()
|
|
85
|
+
loss.backward()
|
|
86
|
+
|
|
87
|
+
for layer in session.record.layers:
|
|
88
|
+
print(layer.full_name, layer.grad_norm, layer.has_zero_grad)
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
Gradient fields stay `None` / `False` if no backward pass runs.
|
|
92
|
+
|
|
93
|
+
## Anomaly checks
|
|
94
|
+
|
|
95
|
+
| Type | Trigger |
|
|
96
|
+
| --- | --- |
|
|
97
|
+
| `nan_activation` | Output tensor contains NaN |
|
|
98
|
+
| `inf_activation` | Output tensor contains Inf |
|
|
99
|
+
| `dead_layer` | Output mean and std both near zero |
|
|
100
|
+
| `exploding_variance` | Std > 100, or >10x spike from previous layer |
|
|
101
|
+
| `high_variance` | Std between 10 and 100 |
|
|
102
|
+
| `empty_output` | No tensor output captured |
|
|
103
|
+
| `zero_gradient` | Backward gradient is all zeros |
|
|
104
|
+
| `nan_gradient` | Backward gradient contains NaN |
|
|
105
|
+
|
|
106
|
+
## JSON format
|
|
107
|
+
|
|
108
|
+
```json
|
|
109
|
+
{
|
|
110
|
+
"metadata": {
|
|
111
|
+
"model_name": "Sequential",
|
|
112
|
+
"pytorch_version": "2.x",
|
|
113
|
+
"python_version": "3.11.x",
|
|
114
|
+
"cuda_available": false,
|
|
115
|
+
"active_device": "cpu",
|
|
116
|
+
"total_time_ms": 1.23
|
|
117
|
+
},
|
|
118
|
+
"layers": [
|
|
119
|
+
{
|
|
120
|
+
"name": "0",
|
|
121
|
+
"type": "Linear",
|
|
122
|
+
"full_name": "0",
|
|
123
|
+
"depth": 1,
|
|
124
|
+
"params": 200960,
|
|
125
|
+
"inputs": [
|
|
126
|
+
{
|
|
127
|
+
"shape": [32, 784],
|
|
128
|
+
"dtype": "torch.float32",
|
|
129
|
+
"device": "cpu",
|
|
130
|
+
"mean": 0.0,
|
|
131
|
+
"std": 1.0,
|
|
132
|
+
"min": -3.0,
|
|
133
|
+
"max": 3.0,
|
|
134
|
+
"nan_count": 0,
|
|
135
|
+
"inf_count": 0
|
|
136
|
+
}
|
|
137
|
+
],
|
|
138
|
+
"outputs": [
|
|
139
|
+
{
|
|
140
|
+
"shape": [32, 256],
|
|
141
|
+
"dtype": "torch.float32",
|
|
142
|
+
"device": "cpu",
|
|
143
|
+
"mean": 0.01,
|
|
144
|
+
"std": 0.58,
|
|
145
|
+
"min": -2.1,
|
|
146
|
+
"max": 2.4,
|
|
147
|
+
"nan_count": 0,
|
|
148
|
+
"inf_count": 0
|
|
149
|
+
}
|
|
150
|
+
],
|
|
151
|
+
"latency_ms": 0.25,
|
|
152
|
+
"grad_norm": null,
|
|
153
|
+
"grad_mean": null,
|
|
154
|
+
"grad_has_nan": false,
|
|
155
|
+
"has_zero_grad": false
|
|
156
|
+
}
|
|
157
|
+
],
|
|
158
|
+
"warnings": []
|
|
159
|
+
}
|
|
160
|
+
```
|
|
161
|
+
|
|
162
|
+
Load a saved trace in Python:
|
|
163
|
+
|
|
164
|
+
```python
|
|
165
|
+
from tracetorch.storage import load_json
|
|
166
|
+
|
|
167
|
+
record = load_json("trace.json")
|
|
168
|
+
print(record.metadata)
|
|
169
|
+
print(record.layers[0].outputs)
|
|
170
|
+
```
|
|
171
|
+
|
|
172
|
+
## API Reference
|
|
173
|
+
|
|
174
|
+
### TraceSession
|
|
175
|
+
|
|
176
|
+
```python
|
|
177
|
+
session = TraceSession(model, capture_stats=True, model_name=None)
|
|
178
|
+
```
|
|
179
|
+
|
|
180
|
+
Context manager for one traced execution window.
|
|
181
|
+
|
|
182
|
+
| Property / Method | Description |
|
|
183
|
+
| --- | --- |
|
|
184
|
+
| `record` | Collected `TraceRecord` |
|
|
185
|
+
| `anomalies` | `list[Anomaly]` populated after exit |
|
|
186
|
+
| `export(path)` | Write JSON, return `Path` |
|
|
187
|
+
| `summary()` | Compact text summary |
|
|
188
|
+
|
|
189
|
+
### Data structures
|
|
190
|
+
|
|
191
|
+
```python
|
|
192
|
+
from tracetorch import TraceRecord, LayerTrace, TensorStats, Anomaly
|
|
193
|
+
```
|
|
194
|
+
|
|
195
|
+
All dataclasses serialize via `.to_dict()`. JSON format is flat for dashboard consumption.
|
|
196
|
+
|
|
197
|
+
### CLI
|
|
198
|
+
|
|
199
|
+
```bash
|
|
200
|
+
tracetorch inspect trace.json
|
|
201
|
+
```
|
|
202
|
+
|
|
203
|
+
Prints model summary, nested layer tree with tensor shapes/params/latency/std, and anomaly warnings.
|
|
204
|
+
|
|
205
|
+
## Architecture
|
|
206
|
+
|
|
207
|
+
```
|
|
208
|
+
PyTorch Model
|
|
209
|
+
│
|
|
210
|
+
PyTorch Hooks
|
|
211
|
+
(forward_pre, forward, backward)
|
|
212
|
+
│
|
|
213
|
+
Trace Collector
|
|
214
|
+
capture shapes, stats, grads
|
|
215
|
+
│
|
|
216
|
+
Anomaly Detector
|
|
217
|
+
8 built-in checks
|
|
218
|
+
│
|
|
219
|
+
Storage Layer
|
|
220
|
+
JSON export / import
|
|
221
|
+
│
|
|
222
|
+
CLI Inspector
|
|
223
|
+
rich terminal output
|
|
224
|
+
```
|
|
225
|
+
|
|
226
|
+
## Development
|
|
227
|
+
|
|
228
|
+
```bash
|
|
229
|
+
git clone https://github.com/phraakture/tracetorch.git
|
|
230
|
+
cd tracetorch
|
|
231
|
+
pip install -e ".[dev]"
|
|
232
|
+
|
|
233
|
+
pytest # tests
|
|
234
|
+
ruff check tracetorch tests # lint
|
|
235
|
+
mypy tracetorch # types
|
|
236
|
+
```
|
|
237
|
+
|
|
238
|
+
Run the example:
|
|
239
|
+
|
|
240
|
+
```bash
|
|
241
|
+
python example.py
|
|
242
|
+
tracetorch inspect trace.json
|
|
243
|
+
```
|
|
244
|
+
|
|
245
|
+
## License
|
|
246
|
+
|
|
247
|
+
MIT
|
|
Binary file
|
|
@@ -0,0 +1,42 @@
|
|
|
1
|
+
"""TraceTorch example — trace a simple transformer model."""
|
|
2
|
+
|
|
3
|
+
import torch
|
|
4
|
+
import torch.nn as nn
|
|
5
|
+
|
|
6
|
+
from tracetorch import TraceSession
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class SimpleTransformer(nn.Module):
|
|
10
|
+
"""A minimal transformer for demonstration."""
|
|
11
|
+
|
|
12
|
+
def __init__(self, vocab_size: int = 1000, d_model: int = 768, n_heads: int = 8) -> None:
|
|
13
|
+
super().__init__()
|
|
14
|
+
self.embedding = nn.Embedding(vocab_size, d_model)
|
|
15
|
+
self.attention = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
|
|
16
|
+
self.linear = nn.Linear(d_model, 10)
|
|
17
|
+
|
|
18
|
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
19
|
+
x = self.embedding(x)
|
|
20
|
+
x, _ = self.attention(x, x, x)
|
|
21
|
+
x = x.mean(dim=1)
|
|
22
|
+
x = self.linear(x)
|
|
23
|
+
return x
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def main() -> None:
|
|
27
|
+
model = SimpleTransformer()
|
|
28
|
+
session = TraceSession(model)
|
|
29
|
+
|
|
30
|
+
with session:
|
|
31
|
+
input_ids = torch.randint(0, 1000, (8, 512))
|
|
32
|
+
model(input_ids)
|
|
33
|
+
|
|
34
|
+
print(session.summary())
|
|
35
|
+
print()
|
|
36
|
+
|
|
37
|
+
session.export("./trace.json")
|
|
38
|
+
print("Trace exported to ./trace.json")
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
if __name__ == "__main__":
|
|
42
|
+
main()
|