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)
),
)
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!