diff --git a/nnodely/exporter/export.py b/nnodely/exporter/export.py index fbfa1ce2..536c3d7c 100644 --- a/nnodely/exporter/export.py +++ b/nnodely/exporter/export.py @@ -1,4 +1,10 @@ -import sys, os, torch, importlib, json +import ast +import importlib +import json +import os +import sys + +import torch from torch.fx import symbolic_trace @@ -37,6 +43,40 @@ def load_model(model_path): model_def = json.load(file) return model_def +def _format_torchscript_dict(dict_node): + if not isinstance(dict_node, ast.Dict): + return None + items = [] + for key, value in zip(dict_node.keys, dict_node.values): + if not isinstance(key, ast.Constant) or not isinstance(key.value, str): + return None + items.append(f"{key.value!r}: {ast.unparse(value)}") + return "{" + ", ".join(items) + "}" + +def _format_torchscript_return(line): + stripped = line.strip() + if not stripped.startswith("return "): + return None + try: + return_node = ast.parse(stripped, mode='exec').body[0] + except SyntaxError: + return None + if not isinstance(return_node, ast.Return) or not isinstance(return_node.value, ast.Tuple): + return None + if len(return_node.value.elts) != 4: + return None + + dicts = [_format_torchscript_dict(dict_node) for dict_node in return_node.value.elts] + if any(dict_str is None for dict_str in dicts): + return None + + names = ["output_dict", "minimize_dict", "closed_loop_update_dict", "connect_update_dict"] + formatted = [] + for name, dict_str in zip(names, dicts): + formatted.append(f" {name} = torch.jit.annotate(Dict[str, torch.Tensor], {dict_str})") + formatted.append(f" return {', '.join(names)}") + return "\n".join(formatted) + def export_python_model(model_def, model, model_path): package_name = __package__.split('.')[0] @@ -51,6 +91,7 @@ def export_python_model(model_def, model, model_path): saved_functions = [] with open(model_path, 'w') as file: + file.write("from typing import Dict, Tuple\n\n") file.write("import torch\n\n") ## write the connect wrap function @@ -155,11 +196,16 @@ def export_python_model(model_def, model, model_path): file.write(" self.all_parameters = torch.nn.ParameterDict(self.all_parameters)\n") file.write(" self.all_constants = torch.nn.ParameterDict(self.all_constants)\n\n") - file.write(" def update(self, closed_loop={}, connect={}, disconnect=False):\n") + file.write(" @torch.jit.ignore\n") + file.write(" def update(self, closed_loop=None, connect=None, disconnect=False):\n") file.write(" pass\n") for line in trace.code.split("\n")[len(saved_functions) + 2:]: - if 'self.relation_forward' in line: + if line.strip() == "def forward(self, kwargs):": + file.write(" def forward(self, kwargs: Dict[str, torch.Tensor]) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:\n") + elif formatted_return := _format_torchscript_return(line): + file.write(f"{formatted_return}\n") + elif 'self.relation_forward' in line: if 'Part' in line or 'Select' in line: attribute = [x for x in line.split() if 'self.relation_forward' in x][0].split('.')[2] old_line = f"self.relation_forward.{attribute}.W" @@ -261,14 +307,14 @@ def export_pythononnx_model(model_def, model_path, model_onnx_path, input_order= file_content = [] with open(model_path, 'r') as file: for line in file: - if 'return ({' in line: - file_content.append(line) - break file_content.append(line) + if 'return ({' in line or 'return output_dict, minimize_dict, closed_loop_update_dict, connect_update_dict' in line: + break file_content = ''.join(file_content) # Replace the forward header file_content = file_content.replace('def forward(self, kwargs):', forward) + file_content = file_content.replace('def forward(self, kwargs: Dict[str, torch.Tensor]) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:', forward) # Perform the substitution for key, value in trace_mapping_input.items(): file_content = file_content.replace(key, value) @@ -403,4 +449,4 @@ def onnx_inference(inputs, path, optimize_graph=False): for item in session.get_inputs(): input_data[item.name] = inputs[item.name] # Run inference - return session.run(output_data, input_data) \ No newline at end of file + return session.run(output_data, input_data) diff --git a/tests/test_export_native_torchscript.py b/tests/test_export_native_torchscript.py new file mode 100644 index 00000000..ec8b7766 --- /dev/null +++ b/tests/test_export_native_torchscript.py @@ -0,0 +1,165 @@ +import importlib.util +import os +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import Dict, Tuple +import unittest + +import torch +from torch import Tensor, nn + + +def deterministic_initializer(indexes, params_size, dict_param): + scale = float(dict_param.get("scale", 0.03)) + value = 0 + for axis, index in enumerate(indexes): + value += (axis + 3) * (index + 5) * 1103515245 + normalized = float((value + 12345) % 4096) / 2048.0 - 1.0 + return scale * normalized + + +def zero_initializer(indexes, params_size, dict_param): + return 0.0 + + +def probe_function(history_signal, current_signal, gain, bias): + return gain * torch.tanh(history_signal + current_signal) + bias + + +def build_probe_model(history_len: int, hidden_dim: int, sample_time: float): + from nnodely import Concatenate, Constant, Fir, Input, Linear, Modely, Output, Parameter, ParamFun, Relu, Select, Tanh, clearNames + + clearNames() + measurement = Input("measurement") + command = Input("command") + encoded_measurement = Relu( + Fir( + hidden_dim, + b=True, + W_init=deterministic_initializer, + W_init_params={"scale": 0.05}, + b_init=zero_initializer, + b_init_params={}, + )(measurement.sw(history_len)) + ) + encoded_command = Relu( + Linear( + hidden_dim, + b=True, + W_init=deterministic_initializer, + W_init_params={"scale": 0.04}, + b_init=zero_initializer, + b_init_params={}, + )(command.last()) + ) + features = Concatenate(encoded_measurement, encoded_command) + hidden = Relu( + Linear( + hidden_dim, + b=True, + W_init=deterministic_initializer, + W_init_params={"scale": 0.03}, + b_init=zero_initializer, + b_init_params={}, + )(features) + ) + raw = Linear( + 2, + b=True, + W_init=deterministic_initializer, + W_init_params={"scale": 0.02}, + b_init=zero_initializer, + b_init_params={}, + )(hidden) + bounded = Tanh(Select(raw, 0)) + residual = ParamFun( + probe_function, + parameters_and_constants=[ + Parameter("probe_gain", values=0.7), + Constant("probe_bias", values=0.05), + ], + )(bounded, Select(raw, 1)) + model = Modely(visualizer=None) + model.addModel( + "native_torchscript_probe", + [ + Output("prediction", residual), + Output("auxiliary", Select(raw, 1)), + ], + ) + model.neuralizeModel(sample_time=sample_time) + return model + + +def export_probe_model(export_dir: Path, name: str, history_len: int, hidden_dim: int, sample_time: float) -> Path: + model_dir = export_dir / name + model_dir.mkdir(parents=True, exist_ok=True) + model = build_probe_model(history_len, hidden_dim, sample_time) + model.exportPythonModel(name=name, model_folder=str(model_dir)) + return model_dir / f"{name}.py" + + +def load_tracer_model(path: Path) -> nn.Module: + module_name = f"nnodely_native_torchscript_probe_{os.getpid()}_{abs(hash(str(path)))}" + spec = importlib.util.spec_from_file_location(module_name, path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load exported model from {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.TracerModel() + + +def make_inputs(batch_size: int, history_len: int, requires_grad: bool) -> Dict[str, Tensor]: + measurement = torch.randn(batch_size, history_len, 1) + command = torch.randn(batch_size, history_len, 1) + if requires_grad: + measurement.requires_grad_() + command.requires_grad_() + return {"measurement": measurement, "command": command} + + +def output_dict(result: Tuple[Dict[str, Tensor], Dict[str, Tensor], Dict[str, Tensor], Dict[str, Tensor]]) -> Dict[str, Tensor]: + return result[0] + + +def backward_probe(module: nn.Module, inputs: Dict[str, Tensor]) -> float: + module.zero_grad(set_to_none=True) + outputs = output_dict(module(inputs)) + loss = outputs["prediction"].pow(2).mean() + 0.1 * outputs["auxiliary"].pow(2).mean() + loss.backward() + if not any(parameter.grad is not None for parameter in module.parameters()): + raise RuntimeError("no parameter gradients found") + if not all(tensor.grad is not None for tensor in inputs.values()): + raise RuntimeError("missing input gradient") + return float(loss.detach()) + + +def max_output_difference(eager: nn.Module, scripted: nn.Module, inputs: Dict[str, Tensor]) -> float: + with torch.no_grad(): + eager_outputs = output_dict(eager(inputs)) + scripted_outputs = output_dict(scripted(inputs)) + return max(float((value - scripted_outputs[key]).abs().max()) for key, value in eager_outputs.items()) + + +class NativeTorchScriptExportTest(unittest.TestCase): + def test_export_python_model_generates_native_torchscript(self): + with TemporaryDirectory() as tmp_dir: + tmp_path = Path(tmp_dir) + exported_model = export_probe_model(tmp_path / "export", "native_torchscript_probe", 5, 8, 0.05) + tracer = load_tracer_model(exported_model) + + eager_loss = backward_probe(tracer.train(), make_inputs(3, 5, True)) + scripted = torch.jit.script(tracer) + scripted_loss = backward_probe(scripted.train(), make_inputs(3, 5, True)) + + comparison_inputs = make_inputs(3, 5, False) + self.assertLessEqual(max_output_difference(tracer.eval(), scripted.eval(), comparison_inputs), 1e-6) + artifact = tmp_path / "native_torchscript_probe.pt" + torch.jit.save(scripted, str(artifact)) + loaded = torch.jit.load(str(artifact), map_location="cpu") + self.assertLessEqual(max_output_difference(tracer.eval(), loaded.eval(), comparison_inputs), 1e-6) + self.assertTrue(torch.isfinite(torch.tensor([eager_loss, scripted_loss])).all()) + + +if __name__ == "__main__": + unittest.main()