Have a question? Connect with the community at the TensorFlow Forum Visit Forum


Loads data and retrains the model based on data for image classification.

Used in the notebooks

Used in the tutorials

train_data Training data.
model_spec Specification for the model.
validation_data Validation data. If None, skips validation process.
batch_size Number of samples per training step. If use_hub_library is False, it represents the base learning rate when train batch size is 256 and it's linear to the batch size.
epochs Number of epochs for training.
train_whole_model If true, the Hub module is trained together with the classification layer on top. Otherwise, only train the top classification layer.
dropout_rate The rate for dropout.
learning_rate Base learning rate when train batch size is 256. Linear to the batch size.
momentum a Python float forwarded to the optimizer. Only used when use_hub_library is True.
shuffle Whether the data should be shuffled.
use_augmentation Use data augmentation for preprocessing.
use_hub_library Use make_image_classifier_lib from tensorflow hub to retrain the model.
warmup_steps Number of warmup steps for warmup schedule on learning rate. If None, the default warmup_steps is used which is the total training steps in two epochs. Only used when use_hub_library is False.
model_dir The location of the model checkpoint files. Only used when use_hub_library is False.
do_train Whether to run training.

An instance based on ImageClassifier.