Skip to content

Model Loading

forge.frontend

The loaders for Forge are light wrappers around TVM's frontend ingestors. Each framework has its own ingestor with varying type-signatures and function parameters that are unique to each framework-ingestor.

Incorrect "Defaults"

The "Defaults" column of each function's parameters in the documentation is incorrect in asserting that all parameters are required. Please reference the type-hints and descriptions for correct details.


PyTorch

Automatic ONNX

If the PyTorch loading fails, Forge will automatically attempt an ONNX ingestion.

forge.from_torch(*args, **kwargs) -> IRModule

Load PyTorch model in the form of a scripted PyTorch model and convert into relay. The companion parameters will be handled automatically.

Parameters:

Name Type Description Default
script_module TopLevelTracedModule object

TorchScripted PyTorch graph Note: We currently only support traces (ie: torch.jit.trace(model, input))

required
input_infos List of tuples

Can be (input name, input shape) or (input name, (input shape, input types)) Graph level input shape and type list The same input names need to be used for deployment, so choose easy to remember names (such as: input0, input1) e.g. [('input0', (1, 2)), ('input1', (3, 4))] or [('input0', ((1, 2), 'int')), ('input1', ((3, 4), 'float'))]

required
custom_convert_map Dictionary of str to Relay op

A custom op conversion map in the same format as _convert_map above

required
default_type str

The default dtype to use when type information is not provided by PyTorch.

required
use_parser_friendly_name bool

When True, replace '.' with `_' in a original parameter name. The Relay text parser treats a variable name followed by a period as a tuple element access, so a variable name like "dense.weight" cannot be parsed correctly. Use this option when you want to run the AnnotateSpans pass on the imported module.

required
keep_quantized_weight bool

Return quantized weights and bias, rather than float ones. PyTorch stores quantized weights in a custom format, so we cannot directly access 8 bit weights as Numpy arrays. We use a PyTorch function to unpack quantized weights into float32 arrays and quantization parameters. By default, we return float32 weights and rely on the QNN lowering and the Relay constant folding pass to quantize weights at compile time. In BYOC use cases, however, we cannot apply the constant folding pass on a QNN graph. If keep_quantized_weight is True, we quantize weights in the frontend using a function that is equivalent to qnn.op.quantize(...) operating on Numpy arrays.

required

Returns:

Type Description
IRModule

Forge IRModule object


ONNX

forge.from_onnx(*args, **kwargs) -> IRModule

Convert a ONNX model into an equivalent Relay Function.

ONNX graphs are represented as Python Protobuf objects. The companion parameters will be handled automatically. However, the input names from onnx graph is vague, mixing inputs and network weights/bias such as "1", "2"... For convenience, we rename the real input names to "input_0", "input_1"... And renaming parameters to "param_0", "param_1"...

By default, ONNX defines models in terms of dynamic shapes. The ONNX importer retains that dynamism upon import, and the compiler attempts to convert the model into a static shapes at compile time. If this fails, there may still be dynamic operations in the model. Not all TVM kernels currently support dynamic shapes, please file an issue on discuss.tvm.apache.org if you hit an error with dynamic kernels.

Parameters:

Name Type Description Default
model protobuf object

ONNX ModelProto after ONNX v1.1.0

required
shape dict of str to tuple

The input shape to the graph

required
dtype str or dict of str to str

The input types to the graph

required
opset int

Override to autodetected opset. This can be helpful for some testing.

required
freeze_params bool

If this parameter is true, the importer will take any provided onnx input values (weights, shapes, etc) and embed them into the relay model as Constants instead of variables. This allows more aggressive optimizations at compile time and helps in making models static if certain inputs represent attributes relay would traditionally consider compile-time constants.

required
convert_config Optional[Dict[str, Any]]

Default config: use_nt_batch_matmul : bool = True True to convert qualified onnx matmul to nn.batch_matmul strict to NT format (transpose_a=False, transpose_b=True).

required
export_node_renamed_model_path str

Export the node renamed onnx model to the path. Some models do not contain names in their nodes. During the conversion, if names of nodes are empty, new names will be assigned based on their op types. The exported model can be the reference to spans.

required

Returns:

Type Description
IRModule

Forge IRModule object


TensorFlow

forge.from_tensorflow(*args, **kwargs) -> IRModule

Load tensorflow graph which is a python tensorflow graph object into relay. The companion parameters will be handled automatically.

Parameters:

Name Type Description Default
graph GraphDef object

Tensorflow GraphDef

required
layout target layout to be used (Optional

NCHW only supported now to enable NHWC models on GPU.

required
shape Dictionary of input dimensions (Optional

Graph level input shape dictionary.

required
outputs List of output tensor names (Optional

if not specified then the last node is assumed as graph output.

required
convert_config Optional[Dict[str, Any]]

Default config: use_dense : bool = True Ture to convert tf.matmul to nn.dense, else to nn.matmul. The nn.dense op requires the data tensor to be non-transposed and weight tensor to be transposed, may insert extra transpose to the original graph. use_nt_batch_matmul : bool = True True to convert tf.batch_matmul to nn.batch_matmul strict to NT format (transpose_a=False, transpose_b=True).

required

Returns:

Type Description
IRModule

Forge IRModule object


Keras

forge.from_keras(*args, **kwargs) -> IRModule

Convert keras model to relay Function.

Parameters:

Name Type Description Default
model Model or Model

The keras model to be converted.

required
shape dict of str to int list/tuple

Input shapes of the model, optional

required
layout str

One of 'NCHW' or 'NHWC', indicates how data should be arranged in the output model. Default layout is 'NCHW' as it in general performs better across TVM.

required

Returns:

Type Description
IRModule

Forge IRModule object


TFLite

forge.from_tflite(*args, **kwargs) -> IRModule

Convert from tflite model into compatible relay Function.

Parameters:

Name Type Description Default
model Model or Model

tflite.Model or tflite.Model.Model (depending on tflite version)

required
shape_dict dict of str to int list/tuple

Input shapes of the model.

required
dtype_dict dict of str to str

Input types of the model.

required

Returns:

Type Description
IRModule

Forge IRModule object