"""Export trained models to HLS4ML for FPGA synthesis."""
from __future__ import annotations
from pathlib import Path
from typing import Any, Union
from .._logging import get_logger
log = get_logger(__name__)
try:
import hls4ml # noqa: F401
_HAS_HLS4ML = True
except ImportError:
_HAS_HLS4ML = False
def _require_hls4ml() -> None:
if not _HAS_HLS4ML:
raise ImportError(
"hls4ml is required for hardware export. "
"Install with: pip install arcade-readout[hardware]"
)
[docs]
def export_hls4ml(
model: Any,
output_dir: Union[str, Path],
*,
backend: str = "Vivado",
clock_period: float = 5.0,
io_type: str = "io_parallel",
reuse_factor: int = 1,
precision: str = "ap_fixed<16,6>",
) -> Path:
"""Convert a trained PyTorch model to an HLS4ML project.
Args:
model: Trained PyTorch ``nn.Module``.
output_dir: Directory for the generated HLS project.
backend: HLS backend (``"Vivado"`` or ``"Quartus"``).
clock_period: Target clock period in nanoseconds.
io_type: ``"io_parallel"`` or ``"io_stream"``.
reuse_factor: Resource reuse factor (higher = smaller area,
longer latency).
precision: Default fixed-point precision for all layers.
Returns:
Path to the generated HLS project directory.
Raises:
ImportError: If hls4ml is not installed.
ValueError: If no Linear layers are found in the model.
"""
_require_hls4ml()
import torch.nn as nn
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
model.eval()
layers: list[dict[str, Any]] = []
for name, module in model.named_modules():
if isinstance(module, nn.Linear):
layers.append(
{
"class_name": "Dense",
"name": name.replace(".", "_"),
"n_in": module.in_features,
"n_out": module.out_features,
}
)
elif isinstance(module, nn.Conv1d):
layers.append(
{
"class_name": "Conv1D",
"name": name.replace(".", "_"),
"n_in": module.in_channels,
"n_filt": module.out_channels,
"filt_width": module.kernel_size[0],
"stride": module.stride[0],
"padding": module.padding[0],
}
)
elif isinstance(module, (nn.ReLU, nn.GELU, nn.Tanh, nn.Sigmoid)):
act_name = type(module).__name__.lower()
layers.append(
{
"class_name": "Activation",
"name": f"act_{name.replace('.', '_')}",
"activation": act_name,
}
)
if not layers:
raise ValueError("No Linear or Conv1d layers found in model")
first = layers[0]
input_dim = first.get("n_in") or first.get("n_in", 1)
log.info(
"Exporting model (input_dim=%d) to HLS4ML: backend=%s, precision=%s",
input_dim,
backend,
precision,
)
config = hls4ml.utils.config_from_pytorch_model(
model,
input_shape=(1, input_dim),
granularity="name",
backend=backend,
default_precision=precision,
default_reuse_factor=reuse_factor,
)
project_path = output_dir / "hls_project"
hls_model = hls4ml.converters.convert_from_pytorch_model(
model,
input_shape=(1, input_dim),
hls_config=config,
output_dir=str(project_path),
backend=backend,
clock_period=clock_period,
io_type=io_type,
)
hls_model.compile()
log.info("HLS4ML project written to %s", project_path)
return project_path