Skip to content

Commit d22e479

Browse files
committed
Document ONNX prerequisites and add [onnx] extras_require
torch>=2.6 changed its ONNX exporter to require onnxscript, so users hitting the export scripts on a fresh install were getting ModuleNotFoundError. This adds: - setup.py: extras_require={'onnx': [onnx, onnxruntime, onnxscript, onnx-simplifier]} so users can pip install nnunetv2_distillation[onnx]. - docs/Distillation.md: a Prerequisites callout before the ONNX export section listing both the extras install and the manual pip command.
1 parent 68d5850 commit d22e479

2 files changed

Lines changed: 27 additions & 1 deletion

File tree

distillation/setup.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,16 @@
77
"torch>=1.6.0",
88
"nnunetv2>=2.5,<3.0",
99
],
10+
extras_require={
11+
# ONNX export deps. torch>=2.6 needs onnxscript for the new exporter;
12+
# onnx-simplifier is only used when the user passes --simplify / -sim.
13+
"onnx": [
14+
"onnx",
15+
"onnxruntime",
16+
"onnxscript",
17+
"onnx-simplifier",
18+
],
19+
},
1020
entry_points={
1121
'console_scripts': [
1222
'nnUNetv2_distillation_train=fast_nnunet_distillation_train:main',

docs/Distillation.md

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,23 @@ Parameter description:
189189

190190
### 3. 📤 Export ONNX Model
191191

192-
Export the trained student model to ONNX format for fast inference:
192+
Export the trained student model to ONNX format for fast inference.
193+
194+
> **Prerequisites:** ONNX export depends on a few extra packages that are
195+
> not part of the core install. The easiest way is the `[onnx]` extra:
196+
>
197+
> ```bash
198+
> pip install -e '.[onnx]' # quotes needed for zsh
199+
> ```
200+
>
201+
> Or install them manually:
202+
>
203+
> ```bash
204+
> pip install onnx onnxruntime onnxscript onnx-simplifier
205+
> ```
206+
>
207+
> `onnxscript` is required by PyTorch 2.6+ for the new ONNX exporter,
208+
> and `onnx-simplifier` is only needed when you pass `--simplify` / `-sim`.
193209
194210
#### Standard Distillation Model Export
195211

0 commit comments

Comments
 (0)