Commit Graph

3 Commits

Author SHA1 Message Date
Madeesh Kannan
eaf66e7431
Add NVTX ranges to TrainablePipe components (#10965)
* `TrainablePipe`: Add NVTX range decorator

* Annotate `TrainablePipe` subclasses with NVTX ranges

* Export function signature to allow introspection of args in tests

* Revert "Annotate `TrainablePipe` subclasses with NVTX ranges"

This reverts commit d8684f7372.

* Revert "Export function signature to allow introspection of args in tests"

This reverts commit f4405ca3ad.

* Revert "`TrainablePipe`: Add NVTX range decorator"

This reverts commit 26536eb6b8.

* Add `spacy.pipes_with_nvtx_range` pipeline callback

* Show warnings for all missing user-defined pipe functions that need to be annotated
Fix imports, typos

* Rename `DEFAULT_ANNOTATABLE_PIPE_METHODS` to `DEFAULT_NVTX_ANNOTATABLE_PIPE_METHODS`
Reorder import

* Walk model nodes directly whilst applying NVTX ranges
Ignore pipe method wrapper when applying range
2022-06-30 11:28:12 +02:00
github-actions[bot]
b0b115ff39
Auto-format code with black (#9530)
Co-authored-by: explosion-bot <explosion-bot@users.noreply.github.com>
2021-10-22 13:03:10 +02:00
Daniël de Kok
1f05f56433
Add the spacy.models_with_nvtx_range.v1 callback (#9124)
* Add the spacy.models_with_nvtx_range.v1 callback

This callback recursively adds NVTX ranges to the Models in each pipe in
a pipeline.

* Fix create_models_with_nvtx_range type signature

* NVTX range: wrap models of all trainable pipes jointly

This avoids that (sub-)models that are shared between pipes get wrapped
twice.

* NVTX range callback: make color configurable

Add forward_color and backprop_color options to set the color for the
NVTX range.

* Move create_models_with_nvtx_range to spacy.ml

* Update create_models_with_nvtx_range for thinc changes

with_nvtx_range now updates an existing node, rather than returning a
wrapper node. So, we can simply walk over the nodes and update them.

* NVTX: use after_pipeline_creation in example
2021-10-20 11:59:48 +02:00