tlt.models.model.BaseModel

class tlt.models.model.BaseModel(model_name: str, framework: FrameworkType, use_case: UseCaseType)[source]

Abstract base class for a pretrained model that can be used for transfer learning

__init__(model_name: str, framework: FrameworkType, use_case: UseCaseType)[source]

Class constructor

Methods

__init__(model_name, framework, use_case)

Class constructor

benchmark(dataset[, saved_model_dir, ...])

Use Intel Neural Compressor to benchmark the model with the dataset argument.

evaluate(dataset)

Evaluate the model using the specified dataset.

export(output_dir)

Export the serialized model to an output directory

load_from_directory(model_dir)

Load a model from a directory

optimize_graph(output_dir[, overwrite_model])

Performs FP32 graph optimization using the Intel Neural Compressor on the model and writes the inference-optimized model to the output_dir.

quantize(output_dir, dataset[, config, ...])

Performs post training quantization using the Intel Neural Compressor on the model using the dataset.

train(dataset, output_dir[, epochs, ...])

Train the model using the specified dataset

Attributes

framework

Framework with which the model is compatible

learning_rate

Learning rate for the model

model_name

Name of the model

preprocessor

Preprocessor for the model

use_case

Use case (or category) to which the model belongs