Skip to content

Make core modules TorchScript compatible#168

Open
MarcelloCeresini wants to merge 1 commit into
tonegas:mainfrom
MarcelloCeresini:feature/torchscriptable
Open

Make core modules TorchScript compatible#168
MarcelloCeresini wants to merge 1 commit into
tonegas:mainfrom
MarcelloCeresini:feature/torchscriptable

Conversation

@MarcelloCeresini

Copy link
Copy Markdown

Summary

This PR updates exportPythonModel so that the generated TracerModel classes can be scripted with torch.jit.script.

What Changed

nnodely/exporter/export.py

  • Generated Python model files now import Dict and Tuple
  • TracerModel.forward now has a TorchScript compatible type signature
  • Output dictionaries are now assigned with torch.jit.annotate(Dict[str, torch.Tensor], ...) so their types are visible and constant for TorchScript
  • The generated update now uses safe non-mutable defaults None instead of {}, and is also marked with torch.jit.ignore since it's not needed for inference

tests/test_export_native_torchscript.py

  • Builds a representative nnodely model with standard classes
  • Exports it with exportPythonModel
  • Imports the generated Python module
  • Runs eager backward
  • Run torch.jit.script(TracerModel())
  • Run scripted backward
  • Save and reload the scripted model
  • Verify eager vs scripted vs reloaded outputs match

Why

TorchScript exports make nnodely models easier to deploy outside Python, in particular they can be loaded from C++ applications through LibTorch, enabling integration into simulators or embedded systems with a small change to the generated export code.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant