tlt.models.image_classification.image_classification_model.ImageClassificationModel

class tlt.models.image_classification.image_classification_model.ImageClassificationModel(image_size, do_fine_tuning: bool, dropout_layer_rate: int, model_name: str, framework: FrameworkType, use_case: UseCaseType)

Base class to represent a pretrained model for image classification

__init__(image_size, do_fine_tuning: bool, dropout_layer_rate: int, model_name: str, framework: FrameworkType, use_case: UseCaseType)

Class constructor

Methods

__init__(image_size, do_fine_tuning, ...)	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.
predict(input_samples)	Generates predictions for the input samples.
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

do_fine_tuning	When True, the weights in all of the model's layers will be trainable.
dropout_layer_rate	The probability of any one node being dropped when a dropout layer is used
framework	Framework with which the model is compatible
image_size	The fixed image size that the pretrained model expects as input, in pixels with equal width and height
learning_rate	Learning rate for the model
model_name	Name of the model
num_classes	The number of output neurons in the model; equal to the number of classes in the dataset
preprocessor	Preprocessor for the model
use_case	Use case (or category) to which the model belongs