torch-kirigami 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.
- torch_kirigami-0.1.0/.gitignore +202 -0
- torch_kirigami-0.1.0/LICENSE +21 -0
- torch_kirigami-0.1.0/PKG-INFO +227 -0
- torch_kirigami-0.1.0/README.md +208 -0
- torch_kirigami-0.1.0/docs/architecture.md +135 -0
- torch_kirigami-0.1.0/docs/dependency-graph-design.md +551 -0
- torch_kirigami-0.1.0/docs/dev-notes/index.md +14 -0
- torch_kirigami-0.1.0/docs/dev-notes/maintainability-review.md +53 -0
- torch_kirigami-0.1.0/docs/dev-notes/planning-performance.md +339 -0
- torch_kirigami-0.1.0/docs/dev-notes/workflow-validation.md +77 -0
- torch_kirigami-0.1.0/docs/development.md +130 -0
- torch_kirigami-0.1.0/docs/getting-started.md +159 -0
- torch_kirigami-0.1.0/docs/index.md +50 -0
- torch_kirigami-0.1.0/docs/measurement.md +135 -0
- torch_kirigami-0.1.0/docs/model-support.md +196 -0
- torch_kirigami-0.1.0/docs/operator-coverage.md +201 -0
- torch_kirigami-0.1.0/docs/persistence.md +121 -0
- torch_kirigami-0.1.0/docs/pruning-design.md +621 -0
- torch_kirigami-0.1.0/docs/sparse-training.md +269 -0
- torch_kirigami-0.1.0/docs/testing-coverage.md +161 -0
- torch_kirigami-0.1.0/docs/testing.md +128 -0
- torch_kirigami-0.1.0/docs/workflow-methods.md +373 -0
- torch_kirigami-0.1.0/examples/custom_rule.py +57 -0
- torch_kirigami-0.1.0/examples/dependency.py +13 -0
- torch_kirigami-0.1.0/examples/fused_attention.py +131 -0
- torch_kirigami-0.1.0/examples/pruning.py +49 -0
- torch_kirigami-0.1.0/examples/workflows/README.md +688 -0
- torch_kirigami-0.1.0/examples/workflows/bn_sparsity.py +323 -0
- torch_kirigami-0.1.0/examples/workflows/gate_pruning.py +348 -0
- torch_kirigami-0.1.0/examples/workflows/group_sparsity.py +366 -0
- torch_kirigami-0.1.0/examples/workflows/imagenet_data.py +182 -0
- torch_kirigami-0.1.0/examples/workflows/imagenet_models.py +212 -0
- torch_kirigami-0.1.0/examples/workflows/isomorphic_pruning.py +1016 -0
- torch_kirigami-0.1.0/examples/workflows/iterative_pruning.py +295 -0
- torch_kirigami-0.1.0/examples/workflows/model_metrics.py +44 -0
- torch_kirigami-0.1.0/examples/workflows/osscar_pruning.py +1026 -0
- torch_kirigami-0.1.0/examples/workflows/prune_finetune.py +468 -0
- torch_kirigami-0.1.0/examples/workflows/requirements-examples.txt +4 -0
- torch_kirigami-0.1.0/examples/workflows/soft_pruning.py +343 -0
- torch_kirigami-0.1.0/examples/workflows/stability_pruning.py +453 -0
- torch_kirigami-0.1.0/examples/workflows/variance_pruning.py +627 -0
- torch_kirigami-0.1.0/examples/workflows/vit_head_pruning.py +292 -0
- torch_kirigami-0.1.0/pyproject.toml +81 -0
- torch_kirigami-0.1.0/tests/__init__.py +1 -0
- torch_kirigami-0.1.0/tests/architecture/__init__.py +1 -0
- torch_kirigami-0.1.0/tests/architecture/contract_inventory.json +4708 -0
- torch_kirigami-0.1.0/tests/architecture/test_coverage_inventory.py +86 -0
- torch_kirigami-0.1.0/tests/architecture/test_extensions.py +161 -0
- torch_kirigami-0.1.0/tests/architecture/test_imports.py +77 -0
- torch_kirigami-0.1.0/tests/architecture/test_registry.py +88 -0
- torch_kirigami-0.1.0/tests/capture/__init__.py +1 -0
- torch_kirigami-0.1.0/tests/capture/test_bindings.py +61 -0
- torch_kirigami-0.1.0/tests/capture/test_boundaries.py +91 -0
- torch_kirigami-0.1.0/tests/capture/test_default_inputs.py +70 -0
- torch_kirigami-0.1.0/tests/capture/test_hooks.py +71 -0
- torch_kirigami-0.1.0/tests/capture/test_isolation.py +188 -0
- torch_kirigami-0.1.0/tests/capture/test_mutation.py +227 -0
- torch_kirigami-0.1.0/tests/capture/test_provenance.py +47 -0
- torch_kirigami-0.1.0/tests/capture/test_shape_expressions.py +76 -0
- torch_kirigami-0.1.0/tests/conftest.py +28 -0
- torch_kirigami-0.1.0/tests/core/__init__.py +1 -0
- torch_kirigami-0.1.0/tests/core/test_constraint_properties.py +134 -0
- torch_kirigami-0.1.0/tests/core/test_constraints.py +74 -0
- torch_kirigami-0.1.0/tests/core/test_coordinate_properties.py +157 -0
- torch_kirigami-0.1.0/tests/core/test_records.py +99 -0
- torch_kirigami-0.1.0/tests/core/test_relations.py +129 -0
- torch_kirigami-0.1.0/tests/core/test_selection.py +155 -0
- torch_kirigami-0.1.0/tests/core/test_selection_invariants.py +108 -0
- torch_kirigami-0.1.0/tests/core/test_single_region_selection.py +93 -0
- torch_kirigami-0.1.0/tests/graph/__init__.py +1 -0
- torch_kirigami-0.1.0/tests/graph/test_barriers.py +85 -0
- torch_kirigami-0.1.0/tests/graph/test_freshness.py +111 -0
- torch_kirigami-0.1.0/tests/graph/test_limits.py +20 -0
- torch_kirigami-0.1.0/tests/graph/test_propagation.py +49 -0
- torch_kirigami-0.1.0/tests/graph/test_provenance.py +38 -0
- torch_kirigami-0.1.0/tests/graph/test_queries.py +87 -0
- torch_kirigami-0.1.0/tests/graph/test_relation_output_cache.py +83 -0
- torch_kirigami-0.1.0/tests/integration/__init__.py +1 -0
- torch_kirigami-0.1.0/tests/integration/test_attention.py +47 -0
- torch_kirigami-0.1.0/tests/integration/test_backend_layouts.py +488 -0
- torch_kirigami-0.1.0/tests/integration/test_capture_equivalence.py +120 -0
- torch_kirigami-0.1.0/tests/integration/test_compact_networks.py +170 -0
- torch_kirigami-0.1.0/tests/integration/test_configuration_cache.py +89 -0
- torch_kirigami-0.1.0/tests/integration/test_constant_guards.py +379 -0
- torch_kirigami-0.1.0/tests/integration/test_constraint_queries.py +223 -0
- torch_kirigami-0.1.0/tests/integration/test_coordinate_compositions.py +202 -0
- torch_kirigami-0.1.0/tests/integration/test_effect_contracts.py +86 -0
- torch_kirigami-0.1.0/tests/integration/test_execution_contexts.py +375 -0
- torch_kirigami-0.1.0/tests/integration/test_extensions.py +223 -0
- torch_kirigami-0.1.0/tests/integration/test_fingerprint_reads.py +179 -0
- torch_kirigami-0.1.0/tests/integration/test_geometric_metric.py +293 -0
- torch_kirigami-0.1.0/tests/integration/test_graph_query_indexes.py +291 -0
- torch_kirigami-0.1.0/tests/integration/test_greedy_completion.py +378 -0
- torch_kirigami-0.1.0/tests/integration/test_group_scaling.py +154 -0
- torch_kirigami-0.1.0/tests/integration/test_isomorphic_workflow.py +770 -0
- torch_kirigami-0.1.0/tests/integration/test_lifecycle.py +275 -0
- torch_kirigami-0.1.0/tests/integration/test_measurement.py +279 -0
- torch_kirigami-0.1.0/tests/integration/test_metadata_guards.py +265 -0
- torch_kirigami-0.1.0/tests/integration/test_osscar_workflow.py +1018 -0
- torch_kirigami-0.1.0/tests/integration/test_parameter_budget.py +289 -0
- torch_kirigami-0.1.0/tests/integration/test_persistence.py +140 -0
- torch_kirigami-0.1.0/tests/integration/test_planning.py +104 -0
- torch_kirigami-0.1.0/tests/integration/test_planning_diagnostics.py +328 -0
- torch_kirigami-0.1.0/tests/integration/test_portable_queries.py +50 -0
- torch_kirigami-0.1.0/tests/integration/test_pretrained_examples.py +1285 -0
- torch_kirigami-0.1.0/tests/integration/test_python_contracts.py +198 -0
- torch_kirigami-0.1.0/tests/integration/test_review_graph_boundaries.py +186 -0
- torch_kirigami-0.1.0/tests/integration/test_review_measurement_boundaries.py +251 -0
- torch_kirigami-0.1.0/tests/integration/test_review_optimizations.py +234 -0
- torch_kirigami-0.1.0/tests/integration/test_review_state_boundaries.py +403 -0
- torch_kirigami-0.1.0/tests/integration/test_shape_provenance.py +218 -0
- torch_kirigami-0.1.0/tests/integration/test_shared_constraints.py +154 -0
- torch_kirigami-0.1.0/tests/integration/test_sparse_workflows.py +176 -0
- torch_kirigami-0.1.0/tests/integration/test_transactions.py +45 -0
- torch_kirigami-0.1.0/tests/integration/test_variance_operators.py +102 -0
- torch_kirigami-0.1.0/tests/integration/test_variance_workflow.py +492 -0
- torch_kirigami-0.1.0/tests/integration/test_view_guards.py +259 -0
- torch_kirigami-0.1.0/tests/integration/test_vit_head_workflow.py +600 -0
- torch_kirigami-0.1.0/tests/operators/__init__.py +1 -0
- torch_kirigami-0.1.0/tests/operators/test_arguments.py +139 -0
- torch_kirigami-0.1.0/tests/operators/test_attention.py +128 -0
- torch_kirigami-0.1.0/tests/operators/test_blocks.py +35 -0
- torch_kirigami-0.1.0/tests/operators/test_convolution.py +201 -0
- torch_kirigami-0.1.0/tests/operators/test_embedding.py +34 -0
- torch_kirigami-0.1.0/tests/operators/test_indexing.py +312 -0
- torch_kirigami-0.1.0/tests/operators/test_matmul.py +63 -0
- torch_kirigami-0.1.0/tests/operators/test_normalization.py +170 -0
- torch_kirigami-0.1.0/tests/operators/test_partitions.py +63 -0
- torch_kirigami-0.1.0/tests/operators/test_pointwise.py +77 -0
- torch_kirigami-0.1.0/tests/operators/test_reduction.py +180 -0
- torch_kirigami-0.1.0/tests/operators/test_registered_entries.py +188 -0
- torch_kirigami-0.1.0/tests/operators/test_shapes.py +259 -0
- torch_kirigami-0.1.0/tests/operators/test_spatial.py +87 -0
- torch_kirigami-0.1.0/tests/persistence/__init__.py +1 -0
- torch_kirigami-0.1.0/tests/persistence/test_checkpoint_state.py +331 -0
- torch_kirigami-0.1.0/tests/persistence/test_checkpoint_validation.py +293 -0
- torch_kirigami-0.1.0/tests/persistence/test_plan.py +65 -0
- torch_kirigami-0.1.0/tests/persistence/test_process.py +44 -0
- torch_kirigami-0.1.0/tests/persistence/test_registration_kinds.py +27 -0
- torch_kirigami-0.1.0/tests/pruning/__init__.py +1 -0
- torch_kirigami-0.1.0/tests/pruning/test_apply.py +181 -0
- torch_kirigami-0.1.0/tests/pruning/test_budget.py +195 -0
- torch_kirigami-0.1.0/tests/pruning/test_completion_order.py +179 -0
- torch_kirigami-0.1.0/tests/pruning/test_conditional_metrics.py +429 -0
- torch_kirigami-0.1.0/tests/pruning/test_context.py +291 -0
- torch_kirigami-0.1.0/tests/pruning/test_dynamic_strategy.py +276 -0
- torch_kirigami-0.1.0/tests/pruning/test_granularity.py +394 -0
- torch_kirigami-0.1.0/tests/pruning/test_identity_axes.py +70 -0
- torch_kirigami-0.1.0/tests/pruning/test_independent_ranking.py +388 -0
- torch_kirigami-0.1.0/tests/pruning/test_meta_output_cache.py +222 -0
- torch_kirigami-0.1.0/tests/pruning/test_metrics.py +53 -0
- torch_kirigami-0.1.0/tests/pruning/test_precision.py +133 -0
- torch_kirigami-0.1.0/tests/pruning/test_recipes.py +34 -0
- torch_kirigami-0.1.0/tests/pruning/test_records.py +91 -0
- torch_kirigami-0.1.0/tests/pruning/test_requirement_index.py +173 -0
- torch_kirigami-0.1.0/tests/pruning/test_sparse_planning.py +118 -0
- torch_kirigami-0.1.0/tests/pruning/test_strategy.py +202 -0
- torch_kirigami-0.1.0/tests/pruning/test_taylor_statistics.py +545 -0
- torch_kirigami-0.1.0/tests/pruning/test_unified_budget.py +237 -0
- torch_kirigami-0.1.0/tests/sparsity/__init__.py +1 -0
- torch_kirigami-0.1.0/tests/sparsity/test_budget_schedule.py +297 -0
- torch_kirigami-0.1.0/tests/sparsity/test_operations_gates.py +485 -0
- torch_kirigami-0.1.0/tests/sparsity/test_regularizers.py +627 -0
- torch_kirigami-0.1.0/tests/support/__init__.py +1 -0
- torch_kirigami-0.1.0/tests/support/graph_helpers.py +9 -0
- torch_kirigami-0.1.0/tests/support/models.py +62 -0
- torch_kirigami-0.1.0/tests/support/numerics.py +33 -0
- torch_kirigami-0.1.0/tests/support/operator_cases.py +441 -0
- torch_kirigami-0.1.0/tests/support/pruning.py +45 -0
- torch_kirigami-0.1.0/tests/support/workflow_fixtures.py +167 -0
- torch_kirigami-0.1.0/torch_kirigami/__init__.py +90 -0
- torch_kirigami-0.1.0/torch_kirigami/bindings.py +250 -0
- torch_kirigami-0.1.0/torch_kirigami/capture.py +738 -0
- torch_kirigami-0.1.0/torch_kirigami/configuration.py +155 -0
- torch_kirigami-0.1.0/torch_kirigami/contracts.py +661 -0
- torch_kirigami-0.1.0/torch_kirigami/errors.py +21 -0
- torch_kirigami-0.1.0/torch_kirigami/graph.py +1021 -0
- torch_kirigami-0.1.0/torch_kirigami/measurement.py +480 -0
- torch_kirigami-0.1.0/torch_kirigami/operation.py +469 -0
- torch_kirigami-0.1.0/torch_kirigami/operators/__init__.py +21 -0
- torch_kirigami-0.1.0/torch_kirigami/operators/attention.py +285 -0
- torch_kirigami-0.1.0/torch_kirigami/operators/coordinates.py +25 -0
- torch_kirigami-0.1.0/torch_kirigami/operators/defaults.py +226 -0
- torch_kirigami-0.1.0/torch_kirigami/operators/effects.py +56 -0
- torch_kirigami-0.1.0/torch_kirigami/operators/extended.py +422 -0
- torch_kirigami-0.1.0/torch_kirigami/operators/indexing.py +365 -0
- torch_kirigami-0.1.0/torch_kirigami/operators/native.py +654 -0
- torch_kirigami-0.1.0/torch_kirigami/operators/shapes.py +284 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/__init__.py +70 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/axis_domains.py +173 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/budget.py +145 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/candidates.py +217 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/checkpoint.py +501 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/granularity.py +101 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/groups.py +102 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/layouts.py +31 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/metrics.py +486 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/plan.py +135 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/planner.py +301 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/pruner.py +354 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/ranking.py +224 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/recipes.py +105 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/rewrite.py +290 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/serialization.py +92 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/state.py +377 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/strategies.py +524 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/types.py +679 -0
- torch_kirigami-0.1.0/torch_kirigami/pruning/validation.py +441 -0
- torch_kirigami-0.1.0/torch_kirigami/py.typed +0 -0
- torch_kirigami-0.1.0/torch_kirigami/regions.py +58 -0
- torch_kirigami-0.1.0/torch_kirigami/registry.py +109 -0
- torch_kirigami-0.1.0/torch_kirigami/relations.py +481 -0
- torch_kirigami-0.1.0/torch_kirigami/selection.py +726 -0
- torch_kirigami-0.1.0/torch_kirigami/sparsity/__init__.py +34 -0
- torch_kirigami-0.1.0/torch_kirigami/sparsity/budget.py +140 -0
- torch_kirigami-0.1.0/torch_kirigami/sparsity/gates.py +198 -0
- torch_kirigami-0.1.0/torch_kirigami/sparsity/operations.py +146 -0
- torch_kirigami-0.1.0/torch_kirigami/sparsity/regularizers.py +135 -0
- torch_kirigami-0.1.0/torch_kirigami/sparsity/schedules.py +193 -0
- torch_kirigami-0.1.0/torch_kirigami/sparsity/values.py +357 -0
|
@@ -0,0 +1,202 @@
|
|
|
1
|
+
# Created by https://www.toptal.com/developers/gitignore/api/visualstudiocode,python
|
|
2
|
+
# Edit at https://www.toptal.com/developers/gitignore?templates=visualstudiocode,python
|
|
3
|
+
|
|
4
|
+
### Python ###
|
|
5
|
+
# Byte-compiled / optimized / DLL files
|
|
6
|
+
__pycache__/
|
|
7
|
+
*.py[cod]
|
|
8
|
+
*$py.class
|
|
9
|
+
|
|
10
|
+
# C extensions
|
|
11
|
+
*.so
|
|
12
|
+
|
|
13
|
+
# Distribution / packaging
|
|
14
|
+
.Python
|
|
15
|
+
build/
|
|
16
|
+
develop-eggs/
|
|
17
|
+
dist/
|
|
18
|
+
downloads/
|
|
19
|
+
eggs/
|
|
20
|
+
.eggs/
|
|
21
|
+
lib/
|
|
22
|
+
lib64/
|
|
23
|
+
parts/
|
|
24
|
+
sdist/
|
|
25
|
+
var/
|
|
26
|
+
wheels/
|
|
27
|
+
share/python-wheels/
|
|
28
|
+
*.egg-info/
|
|
29
|
+
.installed.cfg
|
|
30
|
+
*.egg
|
|
31
|
+
MANIFEST
|
|
32
|
+
|
|
33
|
+
# PyInstaller
|
|
34
|
+
# Usually these files are written by a python script from a template
|
|
35
|
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
|
36
|
+
*.manifest
|
|
37
|
+
*.spec
|
|
38
|
+
|
|
39
|
+
# Installer logs
|
|
40
|
+
pip-log.txt
|
|
41
|
+
pip-delete-this-directory.txt
|
|
42
|
+
|
|
43
|
+
# Unit test / coverage reports
|
|
44
|
+
htmlcov/
|
|
45
|
+
.tox/
|
|
46
|
+
.nox/
|
|
47
|
+
.coverage
|
|
48
|
+
.coverage.*
|
|
49
|
+
.cache
|
|
50
|
+
nosetests.xml
|
|
51
|
+
coverage.xml
|
|
52
|
+
*.cover
|
|
53
|
+
*.py,cover
|
|
54
|
+
.hypothesis/
|
|
55
|
+
.pytest_cache/
|
|
56
|
+
cover/
|
|
57
|
+
|
|
58
|
+
# Translations
|
|
59
|
+
*.mo
|
|
60
|
+
*.pot
|
|
61
|
+
|
|
62
|
+
# Django stuff:
|
|
63
|
+
*.log
|
|
64
|
+
local_settings.py
|
|
65
|
+
db.sqlite3
|
|
66
|
+
db.sqlite3-journal
|
|
67
|
+
|
|
68
|
+
# Flask stuff:
|
|
69
|
+
instance/
|
|
70
|
+
.webassets-cache
|
|
71
|
+
|
|
72
|
+
# Scrapy stuff:
|
|
73
|
+
.scrapy
|
|
74
|
+
|
|
75
|
+
# Sphinx documentation
|
|
76
|
+
docs/_build/
|
|
77
|
+
|
|
78
|
+
# PyBuilder
|
|
79
|
+
.pybuilder/
|
|
80
|
+
target/
|
|
81
|
+
|
|
82
|
+
# Jupyter Notebook
|
|
83
|
+
.ipynb_checkpoints
|
|
84
|
+
|
|
85
|
+
# IPython
|
|
86
|
+
profile_default/
|
|
87
|
+
ipython_config.py
|
|
88
|
+
|
|
89
|
+
# pyenv
|
|
90
|
+
# For a library or package, you might want to ignore these files since the code is
|
|
91
|
+
# intended to run in multiple environments; otherwise, check them in:
|
|
92
|
+
# .python-version
|
|
93
|
+
|
|
94
|
+
# pipenv
|
|
95
|
+
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
|
96
|
+
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
|
97
|
+
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
|
98
|
+
# install all needed dependencies.
|
|
99
|
+
#Pipfile.lock
|
|
100
|
+
|
|
101
|
+
# poetry
|
|
102
|
+
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
|
103
|
+
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
|
104
|
+
# commonly ignored for libraries.
|
|
105
|
+
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
|
106
|
+
#poetry.lock
|
|
107
|
+
|
|
108
|
+
# pdm
|
|
109
|
+
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
|
110
|
+
#pdm.lock
|
|
111
|
+
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
|
112
|
+
# in version control.
|
|
113
|
+
# https://pdm.fming.dev/#use-with-ide
|
|
114
|
+
.pdm.toml
|
|
115
|
+
|
|
116
|
+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
|
117
|
+
__pypackages__/
|
|
118
|
+
|
|
119
|
+
# Celery stuff
|
|
120
|
+
celerybeat-schedule
|
|
121
|
+
celerybeat.pid
|
|
122
|
+
|
|
123
|
+
# SageMath parsed files
|
|
124
|
+
*.sage.py
|
|
125
|
+
|
|
126
|
+
# Environments
|
|
127
|
+
.env
|
|
128
|
+
.venv
|
|
129
|
+
env/
|
|
130
|
+
venv/
|
|
131
|
+
ENV/
|
|
132
|
+
env.bak/
|
|
133
|
+
venv.bak/
|
|
134
|
+
|
|
135
|
+
# Spyder project settings
|
|
136
|
+
.spyderproject
|
|
137
|
+
.spyproject
|
|
138
|
+
|
|
139
|
+
# Rope project settings
|
|
140
|
+
.ropeproject
|
|
141
|
+
|
|
142
|
+
# mkdocs documentation
|
|
143
|
+
/site
|
|
144
|
+
|
|
145
|
+
# mypy
|
|
146
|
+
.mypy_cache/
|
|
147
|
+
.dmypy.json
|
|
148
|
+
dmypy.json
|
|
149
|
+
|
|
150
|
+
# Pyre type checker
|
|
151
|
+
.pyre/
|
|
152
|
+
|
|
153
|
+
# pytype static type analyzer
|
|
154
|
+
.pytype/
|
|
155
|
+
|
|
156
|
+
# Cython debug symbols
|
|
157
|
+
cython_debug/
|
|
158
|
+
|
|
159
|
+
# PyCharm
|
|
160
|
+
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
|
161
|
+
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
|
162
|
+
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
|
163
|
+
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
|
164
|
+
#.idea/
|
|
165
|
+
|
|
166
|
+
### Python Patch ###
|
|
167
|
+
# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
|
|
168
|
+
poetry.toml
|
|
169
|
+
|
|
170
|
+
# ruff
|
|
171
|
+
.ruff_cache/
|
|
172
|
+
|
|
173
|
+
# LSP config files
|
|
174
|
+
pyrightconfig.json
|
|
175
|
+
|
|
176
|
+
### VisualStudioCode ###
|
|
177
|
+
.vscode/*
|
|
178
|
+
!.vscode/settings.json
|
|
179
|
+
!.vscode/tasks.json
|
|
180
|
+
!.vscode/launch.json
|
|
181
|
+
!.vscode/extensions.json
|
|
182
|
+
!.vscode/*.code-snippets
|
|
183
|
+
|
|
184
|
+
# Local History for Visual Studio Code
|
|
185
|
+
.history/
|
|
186
|
+
|
|
187
|
+
# Built Visual Studio Code Extensions
|
|
188
|
+
*.vsix
|
|
189
|
+
|
|
190
|
+
### VisualStudioCode Patch ###
|
|
191
|
+
# Ignore all local history of files
|
|
192
|
+
.history
|
|
193
|
+
.ionide
|
|
194
|
+
|
|
195
|
+
# End of https://www.toptal.com/developers/gitignore/api/visualstudiocode,python
|
|
196
|
+
|
|
197
|
+
/tmp
|
|
198
|
+
/runs/
|
|
199
|
+
/examples/workflows/runs/
|
|
200
|
+
|
|
201
|
+
# Local chat reviews, research notes, and one-off implementation records
|
|
202
|
+
/.local/chat-notes/
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 zaber
|
|
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,227 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: torch-kirigami
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: FX-based structural dependency analysis for PyTorch
|
|
5
|
+
Project-URL: Repository, https://github.com/ZaberKo/torch-kirigami
|
|
6
|
+
Project-URL: Documentation, https://github.com/ZaberKo/torch-kirigami/blob/main/docs/index.md
|
|
7
|
+
Project-URL: Issues, https://github.com/ZaberKo/torch-kirigami/issues
|
|
8
|
+
License-Expression: MIT
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Keywords: fx,pytorch,structural-pruning
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Programming Language :: Python :: 3
|
|
14
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
15
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
16
|
+
Requires-Python: >=3.10
|
|
17
|
+
Requires-Dist: torch>=2.6
|
|
18
|
+
Description-Content-Type: text/markdown
|
|
19
|
+
|
|
20
|
+
# torch-kirigami
|
|
21
|
+
|
|
22
|
+
**Structural dependency analysis and physical pruning for PyTorch.**
|
|
23
|
+
|
|
24
|
+
torch-kirigami determines which tensor regions must change together when a channel, feature, or attention dimension is removed. It separates dependency analysis from pruning decisions and model mutation, so the same structural model supports manual pruning, automatic selection, sparse training, and custom operators.
|
|
25
|
+
|
|
26
|
+
The library requires **Python 3.10+ and PyTorch 2.6+**. PyTorch is its only runtime dependency.
|
|
27
|
+
|
|
28
|
+
## How it works
|
|
29
|
+
|
|
30
|
+
```mermaid
|
|
31
|
+
flowchart LR
|
|
32
|
+
Model["Model"] --> Graph["Dependency graph"]
|
|
33
|
+
Graph --> Plan["Pruning plan"]
|
|
34
|
+
Plan --> Compact["Compact model"]
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
- **Dependency analysis** captures a model with FX, describes tensor regions and operator relations, and computes the impact of a selection without changing the model.
|
|
38
|
+
- **Pruning** discovers candidates, scores and selects them, validates an executable plan, and commits parameter and module-attribute changes.
|
|
39
|
+
- **Sparse training components** provide scalar regularizers, explicit channel gates, parameter projections, cumulative budgets, and schedules. The caller owns the task loss, optimizer, and training loop.
|
|
40
|
+
- **Measurement** reports parameter counts, supported multiply–accumulate operations, and inference latency, including `torch.compile` execution.
|
|
41
|
+
|
|
42
|
+
See the [architecture overview](docs/architecture.md) for the complete component diagram and dependency boundaries.
|
|
43
|
+
|
|
44
|
+
## Install
|
|
45
|
+
|
|
46
|
+
Install the published package:
|
|
47
|
+
|
|
48
|
+
```bash
|
|
49
|
+
uv venv .venv
|
|
50
|
+
source .venv/bin/activate
|
|
51
|
+
uv pip install --torch-backend=auto torch-kirigami
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
For development from the repository root:
|
|
55
|
+
|
|
56
|
+
```bash
|
|
57
|
+
uv venv .venv
|
|
58
|
+
source .venv/bin/activate
|
|
59
|
+
uv pip install --torch-backend=auto -e .
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
The [ImageNet workflow guide](examples/workflows/README.md) installs the additional dependencies needed by the pretrained-model examples.
|
|
63
|
+
|
|
64
|
+
## One-shot automatic pruning
|
|
65
|
+
|
|
66
|
+
Build the dependency graph, discover candidates, and call `prune()` to select and
|
|
67
|
+
physically remove channels in one round. This example requires CUDA; set `device`
|
|
68
|
+
to `"cpu"` for a CPU run. It needs no dataset or training loop.
|
|
69
|
+
|
|
70
|
+
```python
|
|
71
|
+
import torch
|
|
72
|
+
from torch import nn
|
|
73
|
+
|
|
74
|
+
from torch_kirigami import DependencyGraph
|
|
75
|
+
from torch_kirigami.pruning import Greedy, GroupMagnitude, ParameterBudget, Pruner
|
|
76
|
+
|
|
77
|
+
device = "cuda"
|
|
78
|
+
model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3)).to(device).eval()
|
|
79
|
+
x = torch.randn(2, 4, device=device)
|
|
80
|
+
|
|
81
|
+
graph = DependencyGraph.build(model, args=(x,))
|
|
82
|
+
pruner = Pruner(model, graph=graph)
|
|
83
|
+
space = pruner.discover_candidates()
|
|
84
|
+
model, result = pruner.prune(
|
|
85
|
+
space,
|
|
86
|
+
budget=ParameterBudget(max_params=51),
|
|
87
|
+
strategy=Greedy(GroupMagnitude(p=2)),
|
|
88
|
+
)
|
|
89
|
+
|
|
90
|
+
assert model[0].out_features == model[2].in_features == 6
|
|
91
|
+
assert model(x).shape == (2, 3)
|
|
92
|
+
print(result.plan.explain())
|
|
93
|
+
```
|
|
94
|
+
|
|
95
|
+
`prune()` combines `plan()` and `apply()` and modifies the original model in place.
|
|
96
|
+
External input/output dimensions are protected by default. Here the whole model
|
|
97
|
+
shrinks from 67 to 51 parameters, removing two of eight hidden features.
|
|
98
|
+
For another pruning round, rebuild the graph. Create a new optimizer if training
|
|
99
|
+
afterward, because physical pruning replaces parameters.
|
|
100
|
+
|
|
101
|
+
`Greedy` scores once. `DynamicGreedy` rescores after each accepted change using
|
|
102
|
+
the selected joint impact; it costs more and does not run training or refresh
|
|
103
|
+
gradients. Metrics and strategies have explicit extension methods for
|
|
104
|
+
model-specific policies. See the [scoring API](docs/pruning-design.md#scoring-and-strategy-contracts).
|
|
105
|
+
|
|
106
|
+
For a whole-model parameter reduction fraction, use
|
|
107
|
+
`budget=ParameterBudget.from_ratio(model, pruning_ratio=0.05)`. This counts the
|
|
108
|
+
initial parameters and converts the ratio to an absolute cap once. The standalone
|
|
109
|
+
`torch_kirigami.measurement.count_parameters(model)` also exposes that count
|
|
110
|
+
without tracing or executing the model.
|
|
111
|
+
|
|
112
|
+
For a pretrained ResNet one-shot command without fine-tuning, see the
|
|
113
|
+
[workflow guide](examples/workflows/README.md#1-magnitude-or-taylor-pruning-and-fine-tuning).
|
|
114
|
+
|
|
115
|
+
## Inspect a manual pruning plan
|
|
116
|
+
|
|
117
|
+
This small model illustrates the API. For accuracy experiments, use the pretrained ImageNet workflows below.
|
|
118
|
+
|
|
119
|
+
```python
|
|
120
|
+
import torch
|
|
121
|
+
from torch import nn
|
|
122
|
+
|
|
123
|
+
from torch_kirigami import DependencyGraph
|
|
124
|
+
from torch_kirigami.pruning import Pruner
|
|
125
|
+
|
|
126
|
+
model = nn.Sequential(nn.Linear(4, 6), nn.ReLU(), nn.Linear(6, 3))
|
|
127
|
+
x = torch.randn(2, 4)
|
|
128
|
+
|
|
129
|
+
# Capture structure and inspect the dependency closure without mutation.
|
|
130
|
+
graph = DependencyGraph.build(model, args=(x,))
|
|
131
|
+
selection = graph.parameter("0.weight").axis(0).select([1, 4])
|
|
132
|
+
impact = graph.propagate(remove=[selection])
|
|
133
|
+
print(graph.explain(impact))
|
|
134
|
+
|
|
135
|
+
# Validate all physical edits before applying them to the original model.
|
|
136
|
+
pruner = Pruner(model, graph=graph)
|
|
137
|
+
plan = pruner.plan_remove([selection])
|
|
138
|
+
print(plan.explain())
|
|
139
|
+
model, result = pruner.apply(plan)
|
|
140
|
+
|
|
141
|
+
assert model[0].out_features == 4
|
|
142
|
+
assert model[2].in_features == 4
|
|
143
|
+
assert model(x).shape == (2, 3)
|
|
144
|
+
|
|
145
|
+
# Structure changed: rebuild analysis and bind a new optimizer.
|
|
146
|
+
graph = DependencyGraph.build(model, args=(x,))
|
|
147
|
+
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
|
|
148
|
+
```
|
|
149
|
+
|
|
150
|
+
Removing two outputs of the first linear layer also removes their bias entries and the corresponding input columns of the second layer. Public input and output dimensions remain unchanged.
|
|
151
|
+
|
|
152
|
+
For automatic selection, replace the explicit `plan_remove(...)` call with:
|
|
153
|
+
|
|
154
|
+
```python
|
|
155
|
+
from torch_kirigami.pruning import ChannelRatio, Greedy, Magnitude
|
|
156
|
+
|
|
157
|
+
pruner = Pruner(model, graph=graph)
|
|
158
|
+
space = pruner.discover_candidates()
|
|
159
|
+
plan = pruner.plan(space, budget=ChannelRatio(0.25), strategy=Greedy(Magnitude(p=2)))
|
|
160
|
+
```
|
|
161
|
+
|
|
162
|
+
Use a `Pruner` bound to the current graph. `ChannelRatio` converts a reduction
|
|
163
|
+
fraction into upper bounds on final channel widths. `ChannelCount(max_channels,
|
|
164
|
+
channel_axes)` supplies those bounds directly; `ParameterBudget(max_params=...)`
|
|
165
|
+
bounds the final whole-model parameter count. Every budget must be met; an unmet target
|
|
166
|
+
raises `PlanningError` before application. Both use the same dependency and
|
|
167
|
+
execution checks. A resolved impact alone does not guarantee executability or
|
|
168
|
+
numerical equivalence to the unpruned model.
|
|
169
|
+
|
|
170
|
+
## Pretrained ImageNet workflows
|
|
171
|
+
|
|
172
|
+
The [workflow guide](examples/workflows/README.md) provides eleven standalone scripts using torchvision's pretrained **ResNet-18/34/50**, **ViT-B/16/32**, or **ConvNeXt-Tiny**, with ImageNet training and validation kept separate. Basic and iterative pruning target every block's internal widths; the ViT head example adds independently prunable attention heads; Isomorphic uses broader graph-declared candidates; reconstruction methods discover structurally eligible chains. Sparse-training examples declare their BN/gate or regularization scopes. Model choices are method-specific: VBP demonstrates ViT or ConvNeXt MLPs; BN sparsity uses ResNets. The documented commands use full-data defaults with separate training and validation batches of 256, and enable compiled inference for latency measurement:
|
|
173
|
+
|
|
174
|
+
| Workflow | Purpose |
|
|
175
|
+
| --- | --- |
|
|
176
|
+
| Basic pruning | Static/dynamic magnitude or Taylor selection; static FPGM filter-distance criterion |
|
|
177
|
+
| Iterative pruning | Repeated selection toward a final absolute parameter limit |
|
|
178
|
+
| BN sparsity | L1 regularization of ResNet batch-normalization scales |
|
|
179
|
+
| Dependency-group sparsity | Group Lasso or increasing squared-L2 regularization |
|
|
180
|
+
| Soft pruning | Repeated zeroing or gradual norm reduction before physical deletion |
|
|
181
|
+
| Gate pruning | Train explicit channel scales, then prune using their magnitudes |
|
|
182
|
+
| Stability-driven pruning | Monitor retained-channel selections while increasing regularization |
|
|
183
|
+
| Variance-Based Pruning | Calibrate MLP activation variance, prune and compensate the consumer bias; optional fine-tuning |
|
|
184
|
+
| Isomorphic Pruning | Calibrate Taylor gradients, rank within structural families and apply a directly supplied family deletion ratio |
|
|
185
|
+
| OSSCAR | Sequential dense-teacher reconstruction, grouped quadratic deletion, local swaps and weight refitting |
|
|
186
|
+
| ViT heads and FFN | Explicit attention conversion, whole-head and FFN candidates, static group magnitude, fixed residual width |
|
|
187
|
+
|
|
188
|
+
Each workflow reports validation accuracy before and after pruning, optional fine-tuning results, parameter counts, MACs, and latency. These are compact algorithm examples, not reproductions of published benchmark results. [Method selection and adaptations](docs/workflow-methods.md) explains the paper/repository comparisons and the implemented scope.
|
|
189
|
+
|
|
190
|
+
## Documentation
|
|
191
|
+
|
|
192
|
+
Start at the [documentation index](docs/index.md), or choose a path:
|
|
193
|
+
|
|
194
|
+
| Goal | Read |
|
|
195
|
+
| --- | --- |
|
|
196
|
+
| Understand component responsibilities and control flow | [Architecture overview](docs/architecture.md) |
|
|
197
|
+
| Learn the public API step by step | [Getting started](docs/getting-started.md) |
|
|
198
|
+
| Review graph construction, relations, and class contracts | [Dependency graph design](docs/dependency-graph-design.md) |
|
|
199
|
+
| Understand candidate selection and physical execution | [Pruning design](docs/pruning-design.md) |
|
|
200
|
+
| Save plans and restore compact models | [Persistence](docs/persistence.md) |
|
|
201
|
+
| Assemble sparse training and iterative algorithms | [Sparse training](docs/sparse-training.md) |
|
|
202
|
+
| Interpret complexity and latency measurements | [Measurement](docs/measurement.md) |
|
|
203
|
+
| Check supported operators and limitations | [Operator coverage](docs/operator-coverage.md) |
|
|
204
|
+
| Develop, test, or publish a release | [Development guide](docs/development.md) |
|
|
205
|
+
|
|
206
|
+
The executable [custom rule](examples/custom_rule.py) and [fused attention](examples/fused_attention.py) examples demonstrate operator extension.
|
|
207
|
+
|
|
208
|
+
## Operating boundaries
|
|
209
|
+
|
|
210
|
+
Capture uses FX symbolic tracing and shape propagation. Tensor-dependent Python branches, dynamic loops, unknown operators, and edits requiring an unsupported forward rewrite are reported rather than silently approximated. Supported behavior is specific to each operator and pruning axis; consult the coverage guide.
|
|
211
|
+
|
|
212
|
+
A graph is bound to the captured input metadata, structure, module modes, and relevant configuration. Rebuild it after physical pruning or incompatible model changes. Parameter replacement also requires a new optimizer; optimizer state is not migrated automatically.
|
|
213
|
+
|
|
214
|
+
Graph construction isolates example inputs and registered buffers and restores supported RNG state. Model forwards must not mutate parameters or produce external side effects, and capture must not run concurrently with training on the same model.
|
|
215
|
+
|
|
216
|
+
## Development
|
|
217
|
+
|
|
218
|
+
```bash
|
|
219
|
+
uv pip install --torch-backend=auto --group dev -e .
|
|
220
|
+
python -m pytest
|
|
221
|
+
ruff check .
|
|
222
|
+
ruff format --check .
|
|
223
|
+
```
|
|
224
|
+
|
|
225
|
+
Use the activated repository environment. Install dependencies with `uv pip install` and invoke Python and developer tools directly to preserve separately installed workflow packages. The project does not pin a CPU-only PyTorch index. See [development and releases](docs/development.md) for the release process and [testing](docs/testing.md) for optional example dependencies, CUDA checks, and minimum-version validation.
|
|
226
|
+
|
|
227
|
+
Licensed under the [MIT License](LICENSE).
|
|
@@ -0,0 +1,208 @@
|
|
|
1
|
+
# torch-kirigami
|
|
2
|
+
|
|
3
|
+
**Structural dependency analysis and physical pruning for PyTorch.**
|
|
4
|
+
|
|
5
|
+
torch-kirigami determines which tensor regions must change together when a channel, feature, or attention dimension is removed. It separates dependency analysis from pruning decisions and model mutation, so the same structural model supports manual pruning, automatic selection, sparse training, and custom operators.
|
|
6
|
+
|
|
7
|
+
The library requires **Python 3.10+ and PyTorch 2.6+**. PyTorch is its only runtime dependency.
|
|
8
|
+
|
|
9
|
+
## How it works
|
|
10
|
+
|
|
11
|
+
```mermaid
|
|
12
|
+
flowchart LR
|
|
13
|
+
Model["Model"] --> Graph["Dependency graph"]
|
|
14
|
+
Graph --> Plan["Pruning plan"]
|
|
15
|
+
Plan --> Compact["Compact model"]
|
|
16
|
+
```
|
|
17
|
+
|
|
18
|
+
- **Dependency analysis** captures a model with FX, describes tensor regions and operator relations, and computes the impact of a selection without changing the model.
|
|
19
|
+
- **Pruning** discovers candidates, scores and selects them, validates an executable plan, and commits parameter and module-attribute changes.
|
|
20
|
+
- **Sparse training components** provide scalar regularizers, explicit channel gates, parameter projections, cumulative budgets, and schedules. The caller owns the task loss, optimizer, and training loop.
|
|
21
|
+
- **Measurement** reports parameter counts, supported multiply–accumulate operations, and inference latency, including `torch.compile` execution.
|
|
22
|
+
|
|
23
|
+
See the [architecture overview](docs/architecture.md) for the complete component diagram and dependency boundaries.
|
|
24
|
+
|
|
25
|
+
## Install
|
|
26
|
+
|
|
27
|
+
Install the published package:
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
uv venv .venv
|
|
31
|
+
source .venv/bin/activate
|
|
32
|
+
uv pip install --torch-backend=auto torch-kirigami
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
For development from the repository root:
|
|
36
|
+
|
|
37
|
+
```bash
|
|
38
|
+
uv venv .venv
|
|
39
|
+
source .venv/bin/activate
|
|
40
|
+
uv pip install --torch-backend=auto -e .
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
The [ImageNet workflow guide](examples/workflows/README.md) installs the additional dependencies needed by the pretrained-model examples.
|
|
44
|
+
|
|
45
|
+
## One-shot automatic pruning
|
|
46
|
+
|
|
47
|
+
Build the dependency graph, discover candidates, and call `prune()` to select and
|
|
48
|
+
physically remove channels in one round. This example requires CUDA; set `device`
|
|
49
|
+
to `"cpu"` for a CPU run. It needs no dataset or training loop.
|
|
50
|
+
|
|
51
|
+
```python
|
|
52
|
+
import torch
|
|
53
|
+
from torch import nn
|
|
54
|
+
|
|
55
|
+
from torch_kirigami import DependencyGraph
|
|
56
|
+
from torch_kirigami.pruning import Greedy, GroupMagnitude, ParameterBudget, Pruner
|
|
57
|
+
|
|
58
|
+
device = "cuda"
|
|
59
|
+
model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3)).to(device).eval()
|
|
60
|
+
x = torch.randn(2, 4, device=device)
|
|
61
|
+
|
|
62
|
+
graph = DependencyGraph.build(model, args=(x,))
|
|
63
|
+
pruner = Pruner(model, graph=graph)
|
|
64
|
+
space = pruner.discover_candidates()
|
|
65
|
+
model, result = pruner.prune(
|
|
66
|
+
space,
|
|
67
|
+
budget=ParameterBudget(max_params=51),
|
|
68
|
+
strategy=Greedy(GroupMagnitude(p=2)),
|
|
69
|
+
)
|
|
70
|
+
|
|
71
|
+
assert model[0].out_features == model[2].in_features == 6
|
|
72
|
+
assert model(x).shape == (2, 3)
|
|
73
|
+
print(result.plan.explain())
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
`prune()` combines `plan()` and `apply()` and modifies the original model in place.
|
|
77
|
+
External input/output dimensions are protected by default. Here the whole model
|
|
78
|
+
shrinks from 67 to 51 parameters, removing two of eight hidden features.
|
|
79
|
+
For another pruning round, rebuild the graph. Create a new optimizer if training
|
|
80
|
+
afterward, because physical pruning replaces parameters.
|
|
81
|
+
|
|
82
|
+
`Greedy` scores once. `DynamicGreedy` rescores after each accepted change using
|
|
83
|
+
the selected joint impact; it costs more and does not run training or refresh
|
|
84
|
+
gradients. Metrics and strategies have explicit extension methods for
|
|
85
|
+
model-specific policies. See the [scoring API](docs/pruning-design.md#scoring-and-strategy-contracts).
|
|
86
|
+
|
|
87
|
+
For a whole-model parameter reduction fraction, use
|
|
88
|
+
`budget=ParameterBudget.from_ratio(model, pruning_ratio=0.05)`. This counts the
|
|
89
|
+
initial parameters and converts the ratio to an absolute cap once. The standalone
|
|
90
|
+
`torch_kirigami.measurement.count_parameters(model)` also exposes that count
|
|
91
|
+
without tracing or executing the model.
|
|
92
|
+
|
|
93
|
+
For a pretrained ResNet one-shot command without fine-tuning, see the
|
|
94
|
+
[workflow guide](examples/workflows/README.md#1-magnitude-or-taylor-pruning-and-fine-tuning).
|
|
95
|
+
|
|
96
|
+
## Inspect a manual pruning plan
|
|
97
|
+
|
|
98
|
+
This small model illustrates the API. For accuracy experiments, use the pretrained ImageNet workflows below.
|
|
99
|
+
|
|
100
|
+
```python
|
|
101
|
+
import torch
|
|
102
|
+
from torch import nn
|
|
103
|
+
|
|
104
|
+
from torch_kirigami import DependencyGraph
|
|
105
|
+
from torch_kirigami.pruning import Pruner
|
|
106
|
+
|
|
107
|
+
model = nn.Sequential(nn.Linear(4, 6), nn.ReLU(), nn.Linear(6, 3))
|
|
108
|
+
x = torch.randn(2, 4)
|
|
109
|
+
|
|
110
|
+
# Capture structure and inspect the dependency closure without mutation.
|
|
111
|
+
graph = DependencyGraph.build(model, args=(x,))
|
|
112
|
+
selection = graph.parameter("0.weight").axis(0).select([1, 4])
|
|
113
|
+
impact = graph.propagate(remove=[selection])
|
|
114
|
+
print(graph.explain(impact))
|
|
115
|
+
|
|
116
|
+
# Validate all physical edits before applying them to the original model.
|
|
117
|
+
pruner = Pruner(model, graph=graph)
|
|
118
|
+
plan = pruner.plan_remove([selection])
|
|
119
|
+
print(plan.explain())
|
|
120
|
+
model, result = pruner.apply(plan)
|
|
121
|
+
|
|
122
|
+
assert model[0].out_features == 4
|
|
123
|
+
assert model[2].in_features == 4
|
|
124
|
+
assert model(x).shape == (2, 3)
|
|
125
|
+
|
|
126
|
+
# Structure changed: rebuild analysis and bind a new optimizer.
|
|
127
|
+
graph = DependencyGraph.build(model, args=(x,))
|
|
128
|
+
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
|
|
129
|
+
```
|
|
130
|
+
|
|
131
|
+
Removing two outputs of the first linear layer also removes their bias entries and the corresponding input columns of the second layer. Public input and output dimensions remain unchanged.
|
|
132
|
+
|
|
133
|
+
For automatic selection, replace the explicit `plan_remove(...)` call with:
|
|
134
|
+
|
|
135
|
+
```python
|
|
136
|
+
from torch_kirigami.pruning import ChannelRatio, Greedy, Magnitude
|
|
137
|
+
|
|
138
|
+
pruner = Pruner(model, graph=graph)
|
|
139
|
+
space = pruner.discover_candidates()
|
|
140
|
+
plan = pruner.plan(space, budget=ChannelRatio(0.25), strategy=Greedy(Magnitude(p=2)))
|
|
141
|
+
```
|
|
142
|
+
|
|
143
|
+
Use a `Pruner` bound to the current graph. `ChannelRatio` converts a reduction
|
|
144
|
+
fraction into upper bounds on final channel widths. `ChannelCount(max_channels,
|
|
145
|
+
channel_axes)` supplies those bounds directly; `ParameterBudget(max_params=...)`
|
|
146
|
+
bounds the final whole-model parameter count. Every budget must be met; an unmet target
|
|
147
|
+
raises `PlanningError` before application. Both use the same dependency and
|
|
148
|
+
execution checks. A resolved impact alone does not guarantee executability or
|
|
149
|
+
numerical equivalence to the unpruned model.
|
|
150
|
+
|
|
151
|
+
## Pretrained ImageNet workflows
|
|
152
|
+
|
|
153
|
+
The [workflow guide](examples/workflows/README.md) provides eleven standalone scripts using torchvision's pretrained **ResNet-18/34/50**, **ViT-B/16/32**, or **ConvNeXt-Tiny**, with ImageNet training and validation kept separate. Basic and iterative pruning target every block's internal widths; the ViT head example adds independently prunable attention heads; Isomorphic uses broader graph-declared candidates; reconstruction methods discover structurally eligible chains. Sparse-training examples declare their BN/gate or regularization scopes. Model choices are method-specific: VBP demonstrates ViT or ConvNeXt MLPs; BN sparsity uses ResNets. The documented commands use full-data defaults with separate training and validation batches of 256, and enable compiled inference for latency measurement:
|
|
154
|
+
|
|
155
|
+
| Workflow | Purpose |
|
|
156
|
+
| --- | --- |
|
|
157
|
+
| Basic pruning | Static/dynamic magnitude or Taylor selection; static FPGM filter-distance criterion |
|
|
158
|
+
| Iterative pruning | Repeated selection toward a final absolute parameter limit |
|
|
159
|
+
| BN sparsity | L1 regularization of ResNet batch-normalization scales |
|
|
160
|
+
| Dependency-group sparsity | Group Lasso or increasing squared-L2 regularization |
|
|
161
|
+
| Soft pruning | Repeated zeroing or gradual norm reduction before physical deletion |
|
|
162
|
+
| Gate pruning | Train explicit channel scales, then prune using their magnitudes |
|
|
163
|
+
| Stability-driven pruning | Monitor retained-channel selections while increasing regularization |
|
|
164
|
+
| Variance-Based Pruning | Calibrate MLP activation variance, prune and compensate the consumer bias; optional fine-tuning |
|
|
165
|
+
| Isomorphic Pruning | Calibrate Taylor gradients, rank within structural families and apply a directly supplied family deletion ratio |
|
|
166
|
+
| OSSCAR | Sequential dense-teacher reconstruction, grouped quadratic deletion, local swaps and weight refitting |
|
|
167
|
+
| ViT heads and FFN | Explicit attention conversion, whole-head and FFN candidates, static group magnitude, fixed residual width |
|
|
168
|
+
|
|
169
|
+
Each workflow reports validation accuracy before and after pruning, optional fine-tuning results, parameter counts, MACs, and latency. These are compact algorithm examples, not reproductions of published benchmark results. [Method selection and adaptations](docs/workflow-methods.md) explains the paper/repository comparisons and the implemented scope.
|
|
170
|
+
|
|
171
|
+
## Documentation
|
|
172
|
+
|
|
173
|
+
Start at the [documentation index](docs/index.md), or choose a path:
|
|
174
|
+
|
|
175
|
+
| Goal | Read |
|
|
176
|
+
| --- | --- |
|
|
177
|
+
| Understand component responsibilities and control flow | [Architecture overview](docs/architecture.md) |
|
|
178
|
+
| Learn the public API step by step | [Getting started](docs/getting-started.md) |
|
|
179
|
+
| Review graph construction, relations, and class contracts | [Dependency graph design](docs/dependency-graph-design.md) |
|
|
180
|
+
| Understand candidate selection and physical execution | [Pruning design](docs/pruning-design.md) |
|
|
181
|
+
| Save plans and restore compact models | [Persistence](docs/persistence.md) |
|
|
182
|
+
| Assemble sparse training and iterative algorithms | [Sparse training](docs/sparse-training.md) |
|
|
183
|
+
| Interpret complexity and latency measurements | [Measurement](docs/measurement.md) |
|
|
184
|
+
| Check supported operators and limitations | [Operator coverage](docs/operator-coverage.md) |
|
|
185
|
+
| Develop, test, or publish a release | [Development guide](docs/development.md) |
|
|
186
|
+
|
|
187
|
+
The executable [custom rule](examples/custom_rule.py) and [fused attention](examples/fused_attention.py) examples demonstrate operator extension.
|
|
188
|
+
|
|
189
|
+
## Operating boundaries
|
|
190
|
+
|
|
191
|
+
Capture uses FX symbolic tracing and shape propagation. Tensor-dependent Python branches, dynamic loops, unknown operators, and edits requiring an unsupported forward rewrite are reported rather than silently approximated. Supported behavior is specific to each operator and pruning axis; consult the coverage guide.
|
|
192
|
+
|
|
193
|
+
A graph is bound to the captured input metadata, structure, module modes, and relevant configuration. Rebuild it after physical pruning or incompatible model changes. Parameter replacement also requires a new optimizer; optimizer state is not migrated automatically.
|
|
194
|
+
|
|
195
|
+
Graph construction isolates example inputs and registered buffers and restores supported RNG state. Model forwards must not mutate parameters or produce external side effects, and capture must not run concurrently with training on the same model.
|
|
196
|
+
|
|
197
|
+
## Development
|
|
198
|
+
|
|
199
|
+
```bash
|
|
200
|
+
uv pip install --torch-backend=auto --group dev -e .
|
|
201
|
+
python -m pytest
|
|
202
|
+
ruff check .
|
|
203
|
+
ruff format --check .
|
|
204
|
+
```
|
|
205
|
+
|
|
206
|
+
Use the activated repository environment. Install dependencies with `uv pip install` and invoke Python and developer tools directly to preserve separately installed workflow packages. The project does not pin a CPU-only PyTorch index. See [development and releases](docs/development.md) for the release process and [testing](docs/testing.md) for optional example dependencies, CUDA checks, and minimum-version validation.
|
|
207
|
+
|
|
208
|
+
Licensed under the [MIT License](LICENSE).
|