mirror of
https://github.com/explosion/spaCy.git
synced 2024-12-25 17:36:30 +03:00
1f05f56433
* 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
38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
from functools import partial
|
|
from typing import Type, Callable, TYPE_CHECKING
|
|
|
|
from thinc.layers import with_nvtx_range
|
|
from thinc.model import Model, wrap_model_recursive
|
|
|
|
from ..util import registry
|
|
|
|
if TYPE_CHECKING:
|
|
# This lets us add type hints for mypy etc. without causing circular imports
|
|
from ..language import Language # noqa: F401
|
|
|
|
|
|
@registry.callbacks("spacy.models_with_nvtx_range.v1")
|
|
def create_models_with_nvtx_range(
|
|
forward_color: int = -1, backprop_color: int = -1
|
|
) -> Callable[["Language"], "Language"]:
|
|
def models_with_nvtx_range(nlp):
|
|
pipes = [
|
|
pipe
|
|
for _, pipe in nlp.components
|
|
if hasattr(pipe, "is_trainable") and pipe.is_trainable
|
|
]
|
|
|
|
# We need process all models jointly to avoid wrapping callbacks twice.
|
|
models = Model(
|
|
"wrap_with_nvtx_range",
|
|
forward=lambda model, X, is_train: ...,
|
|
layers=[pipe.model for pipe in pipes],
|
|
)
|
|
|
|
for node in models.walk():
|
|
with_nvtx_range(node, forward_color=forward_color, backprop_color=backprop_color)
|
|
|
|
return nlp
|
|
|
|
return models_with_nvtx_range
|