Mojo struct
InputSpec
Specifies a model's input shape and data type (required for TorchScript).
When loading a TorchScript model, you must specify the shape and data type
for each input with an InputSpec
, which you then pass to
InferenceSession.load()
.
For example:
var batch = 1
var seqlen = 128
var input_ids_spec = TensorSpec(DType.int64, batch, seqlen)
var attention_mask_spec = TensorSpec(DType.int64, batch, seqlen)
var session = engine.InferenceSession()
var model = session.load(
"roberta.torchscript",
input_specs=List[InputSpec](
InputSpec(input_ids_spec), InputSpec(attention_mask_spec)
),
)
Implemented traitsβ
AnyType
,
CollectionElement
,
Copyable
,
Movable
Methodsβ
__init__
β
__init__(inout self: Self, spec: TensorSpec)
Create input specifications for one input tensor, as a TensorSpec
. Only applicable for TorchScript models.
Args:
- βspec (
TensorSpec
): Spec for the input. This is the standard libraryTensorSpec
.
__init__(inout self: Self, spec: Optional[List[Optional[SIMD[int64, 1]], 0]], dtype: DType)
Create specifications for one input tensor, as a list of integers. Only applicable for TorchScript models.
If an input supports dynamic shapes, use None
for that dimension
size.
Args:
- βspec (
Optional[List[Optional[SIMD[int64, 1]], 0]]
): Shape of the input, as a list of integers. - βdtype (
DType
): Datatype of the input, from the standard libraryDType
.
__init__(inout self: Self, spec: Optional[List[ShapeElement, 0]], dtype: DType)
Create specifications for one input tensor, as a list of shape elements. Only applicable for TorchScript models.
If an input supports dynamic shapes, use None
or a string dimension
name for that dimension size.
Args:
- βspec (
Optional[List[ShapeElement, 0]]
): Shape of the input, as a list ofShapeElement
values. - βdtype (
DType
): Datatype of the input, from the standard libraryDType
.
__init__(inout self: Self, spec: NoneType, dtype: DType)
Create a specification for a dynamic-rank input. Only applicable for TorchScript models.
Args:
- βspec (
NoneType
): AlwaysNone
. - βdtype (
DType
): Datatype of the input, from the standard libraryDType
.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!
If you'd like to share more information, please report an issue on GitHub
π What went wrong?