tlt.models.image_classification.tf_image_classification_model.TFImageClassificationModel¶
- class tlt.models.image_classification.tf_image_classification_model.TFImageClassificationModel(model_name: str, model=None, optimizer=None, loss=None, **kwargs)[source]¶
Class to represent a TF custom pretrained model for image classification
- __init__(model_name: str, model=None, optimizer=None, loss=None, **kwargs)[source]¶
Class constructor
Methods
__init__
(model_name[, model, optimizer, loss])Class constructor
benchmark
(dataset[, saved_model_dir, ...])Use Intel Neural Compressor to benchmark the model with the dataset argument.
cleanup_saved_objects_for_distributed
()evaluate
(dataset[, use_test_set, callbacks, ...])Evaluate the accuracy of the model on a dataset.
export
(output_dir)Exports a trained model as a saved_model.pb file.
export_for_distributed
([export_dir, ...])Exports the model, optimizer, loss, train data and validation data to the export_dir for distributed script to access.
load_from_directory
(model_dir)Loads a saved model from the specified 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[, return_type, ...])Perform feed-forward inference and predict the classes of the input_samples.
quantize
(output_dir, dataset[, config, ...])Performs post training quantization using the Intel Neural Compressor on the model using the dataset.
set_auto_mixed_precision
(...)Enable auto mixed precision for training.
train
(dataset, output_dir[, epochs, ...])Trains the model using the specified image classification 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