ExTorch.JIT (extorch v0.4.0)

TorchScript model loading, inference, and management.

This module provides functions to load pre-trained TorchScript (.pt) models, run inference, inspect model structure, and manage model lifecycle.

Example

model = ExTorch.JIT.load("model.pt")
input = ExTorch.randn({1, 3, 224, 224})
output = ExTorch.JIT.forward(model, [input])

Summary

Functions

buffers(model)

Get named buffers of a model.

eval(model)

Set a model to evaluation mode.

forward(model, inputs)

Run the forward method on a model.

invoke(model, method_name, inputs \\ [])

Invoke a named method on a model.

load(path, opts \\ [])

Load a TorchScript model from a file.

methods(model)

Get the names of all methods on a model.

modules(model)

Get names of submodules of a model.

parameters(model)

Get named parameters of a model.

save(model, path)

Save a TorchScript model to a file.

to(model, device)

Move a model to a different device.

train(model)

Set a model to training mode.

Functions

buffers(model)

@spec buffers(ExTorch.JIT.Model.t()) :: [{String.t(), ExTorch.Tensor.t()}]

Get named buffers of a model.

Returns

A list of {name, tensor} tuples.

eval(model)

@spec eval(ExTorch.JIT.Model.t()) :: :ok

Set a model to evaluation mode.

forward(model, inputs)

@spec forward(ExTorch.JIT.Model.t(), [ExTorch.Tensor.t()]) :: term()

Run the forward method on a model.

Arguments

model: The loaded model.
inputs: A list of input tensors.

Returns

The model output, which can be a tensor, tuple, list, map, or scalar depending on the model's return type.

Examples

input = ExTorch.randn({1, 784})
output = ExTorch.JIT.forward(model, [input])

invoke(model, method_name, inputs \\ [])

@spec invoke(ExTorch.JIT.Model.t(), String.t(), [ExTorch.Tensor.t()]) :: term()

Invoke a named method on a model.

Arguments

model: The loaded model.
method_name: Name of the method to invoke.
inputs: A list of input tensors.

Returns

The method output.

load(path, opts \\ [])

@spec load(
  String.t(),
  keyword()
) :: ExTorch.JIT.Model.t()

Load a TorchScript model from a file.

Arguments

path: Path to the .pt file.
opts: Keyword list of options.
- :device - Device to load the model onto (default: :cpu).

Returns

A %ExTorch.JIT.Model{} struct.

Examples

model = ExTorch.JIT.load("model.pt")
model = ExTorch.JIT.load("model.pt", device: {:cuda, 0})

methods(model)

@spec methods(ExTorch.JIT.Model.t()) :: [String.t()]

Get the names of all methods on a model.

Returns

A list of method name strings.

modules(model)

@spec modules(ExTorch.JIT.Model.t()) :: [String.t()]

Get names of submodules of a model.

Returns

A list of submodule name strings.

parameters(model)

@spec parameters(ExTorch.JIT.Model.t()) :: [{String.t(), ExTorch.Tensor.t()}]

Get named parameters of a model.

Returns

A list of {name, tensor} tuples.

save(model, path)

@spec save(ExTorch.JIT.Model.t(), String.t()) :: :ok

Save a TorchScript model to a file.

Arguments

model: The model to save.
path: Path to save the .pt file.

to(model, device)

@spec to(ExTorch.JIT.Model.t(), ExTorch.Device.device()) :: ExTorch.JIT.Model.t()

Move a model to a different device.

Arguments

model: The model to move.
device: Target device (e.g., :cpu, {:cuda, 0}).

Returns

A new %ExTorch.JIT.Model{} on the target device.

train(model)

@spec train(ExTorch.JIT.Model.t()) :: :ok

Set a model to training mode.