Export model with post-training transformations for JIT compilation.
This method should be added to model classes that support JIT export.
The model must have _transform_names and _transform_types attributes.
| PARAMETER |
DESCRIPTION |
metadata
|
Additional metadata to include in exported model.
TYPE:
dict | None
DEFAULT:
None
|
| RETURNS |
DESCRIPTION |
ScriptModule
|
JIT-compiled model with all transformations applied.
|
| RAISES |
DESCRIPTION |
AttributeError
|
If model lacks required _transform_names or _transform_types
attributes.
|
Source code in spectre/utils/jit.py
| def jit_export(self, metadata: dict | None = None) -> torch.jit.ScriptModule:
"""
Export model with post-training transformations for JIT compilation.
This method should be added to model classes that support JIT export.
The model must have `_transform_names` and `_transform_types` attributes.
Parameters
----------
metadata : dict | None, optional
Additional metadata to include in exported model.
Returns
-------
torch.jit.ScriptModule
JIT-compiled model with all transformations applied.
Raises
------
AttributeError
If model lacks required `_transform_names` or `_transform_types`
attributes.
"""
# Validate model has required attributes
if not hasattr(self, "_transform_names"):
raise AttributeError(
"Model must have '_transform_names' attribute for JIT export. "
"Ensure model inherits from spectre.core.Model or provides this attribute."
)
if not hasattr(self, "_transform_types"):
raise AttributeError(
"Model must have '_transform_types' attribute for JIT export. "
"Ensure model inherits from spectre.core.Model or provides this attribute."
)
class JITExportableModel(torch.nn.Module):
"""Wrapper for JIT-exportable model with transformations."""
def __init__(self, model: Any, metadata: dict | None = None):
super().__init__()
self.model = model
self._metadata_kv = [f"{k}:{v}" for k, v in (metadata or {}).items()]
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.model(x)
@torch.compile
def describe_transforms(self) -> list[str]:
"""Get description of all transformations."""
return [
f"{n}: {t}"
for n, t in zip(
self.model._transform_names, self.model._transform_types
)
]
@torch.compile
def transform_names(self) -> list[str]:
"""Get names of all transformations."""
return self.model._transform_names.copy()
@torch.compile
def describe_metadata(self) -> list[str]:
"""Get metadata key-value pairs."""
return self._metadata_kv
wrapper = JITExportableModel(self, metadata)
wrapper.eval()
return torch.jit.script(wrapper)
|