spaCy/spacy/ml/__init__.py
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

3 lines
108 B
Python

from .callbacks import create_models_with_nvtx_range # noqa: F401
from .models import * # noqa: F401, F403