nemo_curator.models.audio.sed.tensorrt

TensorRT execution for the CNN14 SED neural core.

Module Contents

Classes

Name	Description
`SedCore`	CNN14 neural core; the checkpoint’s spectrogram frontend stays in PyTorch.
`TensorRTPANNsSEDAdapter`	Run the PANNs CNN14 neural core with a target-specific TensorRT engine.
`TensorRTRunner`	Persistent TensorRT runner with shape-specific context memory.
`TensorRTSed`	Reusable SED adapter preserving the checkpoint’s PyTorch frontend.

Functions

Name	Description
`_sha256`	-
`_trt_dtype_to_torch`	-
`_validate_engine_metadata`	Reject an engine whose immutable build contract differs from this adapter.
`extract_features`	Run the checkpoint’s exact spectrogram and log-mel frontend.
`postprocess`	Restore PANNs framewise output geometry from CNN14 segment outputs.

Data

_TENSORRT_MODEL_TYPE

API

class nemo_curator.models.audio.sed.tensorrt.SedCore(
    model: torch.nn.Module
)

Bases: Module

CNN14 neural core; the checkpoint’s spectrogram frontend stays in PyTorch.

bn0

= model.bn0

conv_block1

= model.conv_block1

conv_block2

= model.conv_block2

conv_block3

= model.conv_block3

conv_block4

= model.conv_block4

conv_block5

= model.conv_block5

conv_block6

= model.conv_block6

fc1

= model.fc1

fc_audioset

= model.fc_audioset

nemo_curator.models.audio.sed.tensorrt.SedCore.forward(
    logmel: torch.Tensor
) -> torch.Tensor

class nemo_curator.models.audio.sed.tensorrt.TensorRTPANNsSEDAdapter(
    checkpoint_path: str | None = None,
    sample_rate: int = _DEFAULT_CHECKPOINT_CONFIG[...,
    model_type: str = _DEFAULT_MODEL_TYPE,
    window_size: int = _DEFAULT_CHECKPOINT_CONFIG[...,
    hop_size: int = _DEFAULT_CHECKPOINT_CONFIG[...,
    mel_bins: int = _DEFAULT_CHECKPOINT_CONFIG[...,
    fmin: int = _DEFAULT_CHECKPOINT_CONFIG[...,
    fmax: int = _DEFAULT_CHECKPOINT_CONFIG[...,
    classes_num: int = _DEFAULT_CHECKPOINT_CONFIG[...,
    pad_short_segments: bool = True,
    tensorrt_engine_path: str | None = None,
    max_duration_sec: float | None = None
)

Dataclass

Bases: PANNsSEDAdapter

Run the PANNs CNN14 neural core with a target-specific TensorRT engine.

Audio preprocessing, checkpoint resolution, batch padding, and the canonical SEDResult contract match PANNsSEDAdapter. The checkpoint’s spectrogram and log-mel frontend remains in PyTorch; only the CNN14 neural core runs in TensorRT. Engines are valid only for the GPU compute capability and TensorRT version recorded in their adjacent JSON sidecar.

_max_input_samples

int | None = field(default=None, init=False, repr=False)

max_duration_sec

float | None = None

tensorrt_engine_path

str | None = None

nemo_curator.models.audio.sed.tensorrt.TensorRTPANNsSEDAdapter.__post_init__() -> None

nemo_curator.models.audio.sed.tensorrt.TensorRTPANNsSEDAdapter.infer_batch(
    items: list[dict[str, object]]
) -> list[nemo_curator.models.audio.sed.base.SEDResult]

Run one TensorRT call and preserve the PANNs adapter result schema.

nemo_curator.models.audio.sed.tensorrt.TensorRTPANNsSEDAdapter.load_model(
    num_gpus: int
) -> None

Load the PyTorch frontend and TensorRT runtime on one CUDA device.

nemo_curator.models.audio.sed.tensorrt.TensorRTPANNsSEDAdapter.unload_model() -> None

Release the TensorRT context, PyTorch frontend, and CUDA cache.

class nemo_curator.models.audio.sed.tensorrt.TensorRTRunner(
    engine_path: str | pathlib.Path,
    expected_metadata: collections.abc.Mapping[str, object]
)

Persistent TensorRT runner with shape-specific context memory.

_context

_device_memory

Tensor | None = None

_engine

_input_names

list[str] = []

_metadata

= self._validate_target(path, expected_metadata)

_output_names

list[str] = []

_runtime

= trt.Runtime(logger)

nemo_curator.models.audio.sed.tensorrt.TensorRTRunner.__call__(
    inputs: torch.Tensor = {}
) -> dict[str, torch.Tensor]

nemo_curator.models.audio.sed.tensorrt.TensorRTRunner._bind_inputs(
    inputs: collections.abc.Mapping[str, torch.Tensor]
) -> torch.device

nemo_curator.models.audio.sed.tensorrt.TensorRTRunner._prepare_device_memory(
    device: torch.device
) -> None

nemo_curator.models.audio.sed.tensorrt.TensorRTRunner._validate_engine_io() -> None

nemo_curator.models.audio.sed.tensorrt.TensorRTRunner._validate_target(
    path: pathlib.Path,
    expected_metadata: collections.abc.Mapping[str, object]
) -> dict[str, object]

nemo_curator.models.audio.sed.tensorrt.TensorRTRunner.close() -> None

class nemo_curator.models.audio.sed.tensorrt.TensorRTSed(
    model: torch.nn.Module,
    engine_path: str | pathlib.Path,
    expected_metadata: collections.abc.Mapping[str, object]
)

Reusable SED adapter preserving the checkpoint’s PyTorch frontend.

logmel

= model.logmel_extractor.to('cuda').eval()

max_input_frames

int

Maximum log-mel frame count accepted by the engine profile.

runner

spectrogram

= model.spectrogram_extractor.to('cuda').eval()

nemo_curator.models.audio.sed.tensorrt.TensorRTSed.__call__(
    waveforms: torch.Tensor
) -> torch.Tensor

Return framewise probabilities for padded [batch, samples] input.

nemo_curator.models.audio.sed.tensorrt.TensorRTSed.close() -> None

nemo_curator.models.audio.sed.tensorrt._sha256(
    path: pathlib.Path
) -> str

nemo_curator.models.audio.sed.tensorrt._trt_dtype_to_torch(
    dtype: object
) -> torch.dtype

nemo_curator.models.audio.sed.tensorrt._validate_engine_metadata(
    engine_path: pathlib.Path,
    expected: collections.abc.Mapping[str, object],
    compute_capability: list[int],
    tensorrt_version: str
) -> dict[str, object]

Reject an engine whose immutable build contract differs from this adapter.

nemo_curator.models.audio.sed.tensorrt.extract_features(
    model: torch.nn.Module,
    waveforms: torch.Tensor
) -> tuple[torch.Tensor, int]

Run the checkpoint’s exact spectrogram and log-mel frontend.

nemo_curator.models.audio.sed.tensorrt.postprocess(
    segmentwise: torch.Tensor,
    frames_num: int
) -> torch.Tensor

Restore PANNs framewise output geometry from CNN14 segment outputs.

nemo_curator.models.audio.sed.tensorrt._TENSORRT_MODEL_TYPE = 'Cnn14_DecisionLevelMax'