|---|---|---|---|
| [View on TensorFlow.org](https://www.tensorflow.org/tfmodels/vision/image_classification) | [Run in Google Colab](https://colab.research.google.com/github/tensorflow/models/blob/master/docs/vision/image_classification.ipynb) | [View on GitHub](https://github.com/tensorflow/models/blob/master/docs/vision/image_classification.ipynb) | [Download notebook](https://storage.googleapis.com/tensorflow_docs/models/docs/vision/image_classification.ipynb) |
This tutorial fine-tunes a Residual Network (ResNet) from the TensorFlow [Model Garden](https://github.com/tensorflow/models) package (`tensorflow-models`) to classify images in the [CIFAR](https://www.cs.toronto.edu/%7Ekriz/cifar.html) dataset.
Model Garden contains a collection of state-of-the-art vision models, implemented with TensorFlow's high-level APIs. The implementations demonstrate the best practices for modeling, letting users to take full advantage of TensorFlow for their research and product development.
This tutorial uses a [ResNet](https://arxiv.org/pdf/1512.03385.pdf) model, a state-of-the-art image classifier. This tutorial uses the ResNet-18 model, a convolutional neural network with 18 layers.
This tutorial demonstrates how to:
1. Use models from the TensorFlow Models package.
2. Fine-tune a pre-built ResNet for image classification.
3. Export the tuned ResNet model.
## Setup
Install and import the necessary modules.
pip install -U -q "tf-models-official"
Import TensorFlow, TensorFlow Datasets, and a few helper libraries.
import pprint
import tempfile
from IPython import display
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_datasets as tfds
The `tensorflow_models` package contains the ResNet vision model, and the `official.vision.serving` model contains the function to save and export the tuned model.
import tensorflow_models as tfm
# These are not in the tfm public API for v2.9. They will be available in v2.10
from official.vision.serving import export_saved_model_lib
import official.core.train_lib
```
/tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tensorflow_hub/__init__.py:61: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.
from pkg_resources import parse_version
```
## Configure the ResNet-18 model for the Cifar-10 dataset
The CIFAR10 dataset contains 60,000 color images in mutually exclusive 10 classes, with 6,000 images in each class.
In Model Garden, the collections of parameters that define a model are called *configs* . Model Garden can create a config based on a known set of parameters via a [factory](https://en.wikipedia.org/wiki/Factory_method_pattern).
Use the `resnet_imagenet` factory configuration, as defined by [`tfm.vision.configs.image_classification.image_classification_imagenet`](https://www.tensorflow.org/api_docs/python/tfm/vision/configs/image_classification/image_classification_imagenet). The configuration is set up to train ResNet to converge on [ImageNet](https://www.image-net.org/).
exp_config = tfm.core.exp_factory.get_exp_config('resnet_imagenet')
tfds_name = 'cifar10'
ds,ds_info = tfds.load(
tfds_name,
with_info=True)
ds_info
```
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1793449324.084852 49795 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 13638 MB memory: -> device: 0, name: Tesla T4, pci bus id: 0000:00:05.0, compute capability: 7.5
I0000 00:00:1793449324.087366 49795 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 13756 MB memory: -> device: 1, name: Tesla T4, pci bus id: 0000:00:06.0, compute capability: 7.5
I0000 00:00:1793449324.089689 49795 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 13756 MB memory: -> device: 2, name: Tesla T4, pci bus id: 0000:00:07.0, compute capability: 7.5
I0000 00:00:1793449324.092038 49795 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 13756 MB memory: -> device: 3, name: Tesla T4, pci bus id: 0000:00:08.0, compute capability: 7.5
tfds.core.DatasetInfo(
name='cifar10',
full_name='cifar10/3.0.2',
description="""
The CIFAR-10 dataset consists of 60000 32x32 colour images in 10 classes, with 6000 images per class. There are 50000 training images and 10000 test images.
""",
homepage='https://www.cs.toronto.edu/~kriz/cifar.html',
data_dir='gs://tensorflow-datasets/datasets/cifar10/3.0.2',
file_format=tfrecord,
download_size=162.17 MiB,
dataset_size=132.40 MiB,
features=FeaturesDict({
'id': Text(shape=(), dtype=string),
'image': Image(shape=(32, 32, 3), dtype=uint8),
'label': ClassLabel(shape=(), dtype=int64, num_classes=10),
}),
supervised_keys=('image', 'label'),
disable_shuffling=False,
nondeterministic_order=False,
splits={
'test': ,
'train': ,
},
citation="""@TECHREPORT{Krizhevsky09learningmultiple,
author = {Alex Krizhevsky},
title = {Learning multiple layers of features from tiny images},
institution = {},
year = {2009}
}""",
)
```
Adjust the model and dataset configurations so that it works with Cifar-10 (`cifar10`).
# Configure model
exp_config.task.model.num_classes = 10
exp_config.task.model.input_size = list(ds_info.features["image"].shape)
exp_config.task.model.backbone.resnet.model_id = 18
# Configure training and testing data
batch_size = 128
exp_config.task.train_data.input_path = ''
exp_config.task.train_data.tfds_name = tfds_name
exp_config.task.train_data.tfds_split = 'train'
exp_config.task.train_data.global_batch_size = batch_size
exp_config.task.validation_data.input_path = ''
exp_config.task.validation_data.tfds_name = tfds_name
exp_config.task.validation_data.tfds_split = 'test'
exp_config.task.validation_data.global_batch_size = batch_size
Adjust the trainer configuration.
logical_device_names = [logical_device.name for logical_device in tf.config.list_logical_devices()]
if 'GPU' in ''.join(logical_device_names):
print('This may be broken in Colab.')
device = 'GPU'
elif 'TPU' in ''.join(logical_device_names):
print('This may be broken in Colab.')
device = 'TPU'
else:
print('Running on CPU is slow, so only train for a few steps.')
device = 'CPU'
if device=='CPU':
train_steps = 20
exp_config.trainer.steps_per_loop = 5
else:
train_steps=5000
exp_config.trainer.steps_per_loop = 100
exp_config.trainer.summary_interval = 100
exp_config.trainer.checkpoint_interval = train_steps
exp_config.trainer.validation_interval = 1000
exp_config.trainer.validation_steps = ds_info.splits['test'].num_examples // batch_size
exp_config.trainer.train_steps = train_steps
exp_config.trainer.optimizer_config.learning_rate.type = 'cosine'
exp_config.trainer.optimizer_config.learning_rate.cosine.decay_steps = train_steps
exp_config.trainer.optimizer_config.learning_rate.cosine.initial_learning_rate = 0.1
exp_config.trainer.optimizer_config.warmup.linear.warmup_steps = 100
```
This may be broken in Colab.
```
Print the modified configuration.
pprint.pprint(exp_config.as_dict())
display.Javascript("google.colab.output.setIframeHeight('300px');")
```
{'runtime': {'all_reduce_alg': None,
'batchnorm_spatial_persistent': False,
'dataset_num_private_threads': None,
'default_shard_dim': -1,
'distribution_strategy': 'mirrored',
'enable_xla': True,
'gpu_thread_mode': None,
'loss_scale': None,
'mixed_precision_dtype': None,
'num_cores_per_replica': 1,
'num_gpus': 0,
'num_packs': 1,
'per_gpu_thread_count': 0,
'run_eagerly': False,
'task_index': -1,
'tpu': None,
'tpu_enable_xla_dynamic_padder': None,
'use_tpu_mp_strategy': False,
'worker_hosts': None},
'task': {'allow_image_summary': False,
'differential_privacy_config': None,
'eval_input_partition_dims': [],
'evaluation': {'precision_and_recall_thresholds': None,
'report_per_class_precision_and_recall': False,
'top_k': 5},
'freeze_backbone': False,
'init_checkpoint': None,
'init_checkpoint_modules': 'all',
'losses': {'l2_weight_decay': 0.0001,
'label_smoothing': 0.0,
'loss_weight': 1.0,
'one_hot': True,
'soft_labels': False,
'use_binary_cross_entropy': False},
'model': {'add_head_batch_norm': False,
'backbone': {'resnet': {'bn_trainable': True,
'depth_multiplier': 1.0,
'model_id': 18,
'replace_stem_max_pool': False,
'resnetd_shortcut': False,
'scale_stem': True,
'se_ratio': 0.0,
'stem_type': 'v0',
'stochastic_depth_drop_rate': 0.0},
'type': 'resnet'},
'dropout_rate': 0.0,
'input_size': [32, 32, 3],
'kernel_initializer': 'random_uniform',
'norm_activation': {'activation': 'relu',
'norm_epsilon': 1e-05,
'norm_momentum': 0.9,
'use_sync_bn': False},
'num_classes': 10,
'output_softmax': False},
'model_output_keys': [],
'name': None,
'train_data': {'apply_tf_data_service_before_batching': False,
'aug_crop': True,
'aug_policy': None,
'aug_rand_hflip': True,
'aug_type': None,
'autotune_algorithm': None,
'block_length': 1,
'cache': False,
'center_crop_fraction': 0.875,
'color_jitter': 0.0,
'crop_area_range': (0.08, 1.0),
'cycle_length': 10,
'decode_jpeg_only': True,
'decoder': {'simple_decoder': {'attribute_names': [],
'mask_binarize_threshold': None,
'regenerate_source_id': False},
'type': 'simple_decoder'},
'deterministic': None,
'drop_remainder': True,
'dtype': 'float32',
'enable_shared_tf_data_service_between_parallel_trainers': False,
'enable_tf_data_service': False,
'file_type': 'tfrecord',
'global_batch_size': 128,
'image_field_key': 'image/encoded',
'input_path': '',
'is_multilabel': False,
'is_training': True,
'label_field_key': 'image/class/label',
'mixup_and_cutmix': None,
'prefetch_buffer_size': None,
'ram_budget': None,
'randaug_magnitude': 10,
'random_erasing': None,
'repeated_augment': None,
'seed': None,
'sharding': True,
'shuffle_buffer_size': 10000,
'tf_data_service_address': None,
'tf_data_service_job_name': None,
'tf_resize_method': 'bilinear',
'tfds_as_supervised': False,
'tfds_data_dir': '',
'tfds_name': 'cifar10',
'tfds_skip_decoding_feature': '',
'tfds_split': 'train',
'three_augment': False,
'trainer_id': None,
'weights': None},
'train_input_partition_dims': [],
'validation_data': {'apply_tf_data_service_before_batching': False,
'aug_crop': True,
'aug_policy': None,
'aug_rand_hflip': True,
'aug_type': None,
'autotune_algorithm': None,
'block_length': 1,
'cache': False,
'center_crop_fraction': 0.875,
'color_jitter': 0.0,
'crop_area_range': (0.08, 1.0),
'cycle_length': 10,
'decode_jpeg_only': True,
'decoder': {'simple_decoder': {'attribute_names': [],
'mask_binarize_threshold': None,
'regenerate_source_id': False},
'type': 'simple_decoder'},
'deterministic': None,
'drop_remainder': True,
'dtype': 'float32',
'enable_shared_tf_data_service_between_parallel_trainers': False,
'enable_tf_data_service': False,
'file_type': 'tfrecord',
'global_batch_size': 128,
'image_field_key': 'image/encoded',
'input_path': '',
'is_multilabel': False,
'is_training': False,
'label_field_key': 'image/class/label',
'mixup_and_cutmix': None,
'prefetch_buffer_size': None,
'ram_budget': None,
'randaug_magnitude': 10,
'random_erasing': None,
'repeated_augment': None,
'seed': None,
'sharding': True,
'shuffle_buffer_size': 10000,
'tf_data_service_address': None,
'tf_data_service_job_name': None,
'tf_resize_method': 'bilinear',
'tfds_as_supervised': False,
'tfds_data_dir': '',
'tfds_name': 'cifar10',
'tfds_skip_decoding_feature': '',
'tfds_split': 'test',
'three_augment': False,
'trainer_id': None,
'weights': None} },
'trainer': {'allow_tpu_summary': False,
'best_checkpoint_eval_metric': '',
'best_checkpoint_export_subdir': '',
'best_checkpoint_metric_comp': 'higher',
'checkpoint_interval': 5000,
'continuous_eval_timeout': 3600,
'eval_tf_function': True,
'eval_tf_while_loop': False,
'loss_upper_bound': 1000000.0,
'max_to_keep': 5,
'optimizer_config': {'ema': None,
'learning_rate': {'cosine': {'alpha': 0.0,
'decay_steps': 5000,
'initial_learning_rate': 0.1,
'name': 'CosineDecay',
'offset': 0},
'type': 'cosine'},
'optimizer': {'sgd': {'clipnorm': None,
'clipvalue': None,
'decay': 0.0,
'global_clipnorm': None,
'momentum': 0.9,
'name': 'SGD',
'nesterov': False},
'type': 'sgd'},
'warmup': {'linear': {'name': 'linear',
'warmup_learning_rate': 0,
'warmup_steps': 100},
'type': 'linear'} },
'preemption_on_demand_checkpoint': True,
'recovery_begin_steps': 0,
'recovery_max_trials': 0,
'steps_per_loop': 100,
'summary_interval': 100,
'train_steps': 5000,
'train_tf_function': True,
'train_tf_while_loop': True,
'validation_interval': 1000,
'validation_steps': 78,
'validation_summary_subdir': 'validation'} }
```
Set up the distribution strategy.
logical_device_names = [logical_device.name for logical_device in tf.config.list_logical_devices()]
if exp_config.runtime.mixed_precision_dtype == tf.float16:
tf.keras.mixed_precision.set_global_policy('mixed_float16')
if 'GPU' in ''.join(logical_device_names):
distribution_strategy = tf.distribute.MirroredStrategy()
elif 'TPU' in ''.join(logical_device_names):
tf.tpu.experimental.initialize_tpu_system()
tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='/device:TPU_SYSTEM:0')
distribution_strategy = tf.distribute.experimental.TPUStrategy(tpu)
else:
print('Warning: this will be really slow.')
distribution_strategy = tf.distribute.OneDeviceStrategy(logical_device_names[0])
```
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
```
Create the `Task` object ([`tfm.core.base_task.Task`](https://www.tensorflow.org/api_docs/python/tfm/core/base_task/Task)) from the [`config_definitions.TaskConfig`](https://www.tensorflow.org/api_docs/python/tfm/core/config_definitions/TaskConfig).
The `Task` object has all the methods necessary for building the dataset, building the model, and running training \& evaluation. These methods are driven by [`tfm.core.train_lib.run_experiment`](https://www.tensorflow.org/api_docs/python/tfm/core/train_lib/run_experiment).
with distribution_strategy.scope():
model_dir = tempfile.mkdtemp()
task = tfm.core.task_factory.get_task(exp_config.task, logging_dir=model_dir)
# tf.keras.utils.plot_model(task.build_model(), show_shapes=True)
for images, labels in task.build_inputs(exp_config.task.train_data).take(1):
print()
print(f'images.shape: {str(images.shape):16} images.dtype: {images.dtype!r}')
print(f'labels.shape: {str(labels.shape):16} labels.dtype: {labels.dtype!r}')
```
images.shape: (128, 32, 32, 3) images.dtype: tf.float32
labels.shape: (128,) labels.dtype: tf.int32
2026-10-31 12:22:07.612783: W tensorflow/core/kernels/data/cache_dataset_ops.cc:917] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.
```
## Visualize the training data
The dataloader applies a z-score normalization using
[`preprocess_ops.normalize_image(image, offset=MEAN_RGB, scale=STDDEV_RGB)`](https://www.tensorflow.org/api_docs/python/tfm/vision/preprocess_ops/normalize_image), so the images returned by the dataset can't be directly displayed by standard tools. The visualization code needs to rescale the data into the \[0,1\] range.
plt.hist(images.numpy().flatten());
Use `ds_info` (which is an instance of [`tfds.core.DatasetInfo`](https://www.tensorflow.org/datasets/api_docs/python/tfds/core/DatasetInfo)) to lookup the text descriptions of each class ID.
label_info = ds_info.features['label']
label_info.int2str(1)
```
'automobile'
```
Visualize a batch of the data.
def show_batch(images, labels, predictions=None):
plt.figure(figsize=(10, 10))
min = images.numpy().min()
max = images.numpy().max()
delta = max - min
for i in range(12):
plt.subplot(6, 6, i + 1)
plt.imshow((images[i]-min) / delta)
if predictions is None:
plt.title(label_info.int2str(labels[i]))
else:
if labels[i] == predictions[i]:
color = 'g'
else:
color = 'r'
plt.title(label_info.int2str(predictions[i]), color=color)
plt.axis("off")
plt.figure(figsize=(10, 10))
for images, labels in task.build_inputs(exp_config.task.train_data).take(1):
show_batch(images, labels)
```
2026-10-31 12:22:09.474945: W tensorflow/core/kernels/data/cache_dataset_ops.cc:917] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.
```
## Visualize the testing data
Visualize a batch of images from the validation dataset.
plt.figure(figsize=(10, 10));
for images, labels in task.build_inputs(exp_config.task.validation_data).take(1):
show_batch(images, labels)
```
2026-10-31 12:22:10.781469: W tensorflow/core/kernels/data/cache_dataset_ops.cc:917] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.
```
## Train and evaluate
model, eval_logs = tfm.core.train_lib.run_experiment(
distribution_strategy=distribution_strategy,
task=task,
mode='train_and_eval',
params=exp_config,
model_dir=model_dir,
run_post_eval=True)
```
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
restoring or initializing model...
INFO:tensorflow:Customized initialization is done through the passed `init_fn`.
INFO:tensorflow:Customized initialization is done through the passed `init_fn`.
train | step: 0 | training until step 1000...
INFO:tensorflow:Collective all_reduce tensors: 65 all_reduces, num_devices = 4, group_size = 4, implementation = CommunicationImplementation.NCCL, num_packs = 1
INFO:tensorflow:Collective all_reduce tensors: 65 all_reduces, num_devices = 4, group_size = 4, implementation = CommunicationImplementation.NCCL, num_packs = 1
INFO:tensorflow:Collective all_reduce tensors: 65 all_reduces, num_devices = 4, group_size = 4, implementation = CommunicationImplementation.NCCL, num_packs = 1
INFO:tensorflow:Collective all_reduce tensors: 65 all_reduces, num_devices = 4, group_size = 4, implementation = CommunicationImplementation.NCCL, num_packs = 1
train | step: 100 | steps/sec: 4.8 | output:
{'accuracy': np.float32(0.19898437),
'learning_rate': np.float32(0.09990134),
'top_5_accuracy': np.float32(0.699375),
'training_loss': np.float32(2.7429802)}
saved checkpoint to /tmpfs/tmp/tmp9ai4tfkc/ckpt-100.
train | step: 200 | steps/sec: 22.0 | output:
{'accuracy': np.float32(0.24640626),
'learning_rate': np.float32(0.09960574),
'top_5_accuracy': np.float32(0.778125),
'training_loss': np.float32(2.5309942)}
train | step: 300 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.30484375),
'learning_rate': np.float32(0.09911436),
'top_5_accuracy': np.float32(0.8309375),
'training_loss': np.float32(2.1693156)}
train | step: 400 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.3284375),
'learning_rate': np.float32(0.09842916),
'top_5_accuracy': np.float32(0.8409375),
'training_loss': np.float32(2.0995054)}
train | step: 500 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.3413281),
'learning_rate': np.float32(0.09755283),
'top_5_accuracy': np.float32(0.84171873),
'training_loss': np.float32(2.0678263)}
train | step: 600 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.35070312),
'learning_rate': np.float32(0.096488826),
'top_5_accuracy': np.float32(0.8509375),
'training_loss': np.float32(2.020701)}
train | step: 700 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.3690625),
'learning_rate': np.float32(0.09524136),
'top_5_accuracy': np.float32(0.8621875),
'training_loss': np.float32(1.9820873)}
train | step: 800 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.3753125),
'learning_rate': np.float32(0.09381534),
'top_5_accuracy': np.float32(0.8659375),
'training_loss': np.float32(1.9574311)}
train | step: 900 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.3905469),
'learning_rate': np.float32(0.092216395),
'top_5_accuracy': np.float32(0.8726562),
'training_loss': np.float32(1.923374)}
train | step: 1000 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.41179687),
'learning_rate': np.float32(0.090450846),
'top_5_accuracy': np.float32(0.8815625),
'training_loss': np.float32(1.8723639)}
eval | step: 1000 | running 78 steps of evaluation...
eval | step: 1000 | steps/sec: 22.6 | eval time: 3.4 sec | output:
{'accuracy': np.float32(0.5051082),
'steps_per_second': 22.60920626105475,
'top_5_accuracy': np.float32(0.9289864),
'validation_loss': np.float32(1.6148832)}
train | step: 1000 | training until step 2000...
train | step: 1100 | steps/sec: 14.4 | output:
{'accuracy': np.float32(0.4242969),
'learning_rate': np.float32(0.08852567),
'top_5_accuracy': np.float32(0.8779687),
'training_loss': np.float32(1.8477825)}
train | step: 1200 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.4259375),
'learning_rate': np.float32(0.08644843),
'top_5_accuracy': np.float32(0.8867969),
'training_loss': np.float32(1.8346292)}
train | step: 1300 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.43976563),
'learning_rate': np.float32(0.08422736),
'top_5_accuracy': np.float32(0.8960937),
'training_loss': np.float32(1.7870352)}
train | step: 1400 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.43273437),
'learning_rate': np.float32(0.081871204),
'top_5_accuracy': np.float32(0.88773435),
'training_loss': np.float32(1.8033347)}
train | step: 1500 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.44515625),
'learning_rate': np.float32(0.07938927),
'top_5_accuracy': np.float32(0.895625),
'training_loss': np.float32(1.7762802)}
train | step: 1600 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.45398438),
'learning_rate': np.float32(0.07679134),
'top_5_accuracy': np.float32(0.90171874),
'training_loss': np.float32(1.74652)}
train | step: 1700 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.45679688),
'learning_rate': np.float32(0.07408768),
'top_5_accuracy': np.float32(0.9027344),
'training_loss': np.float32(1.7297899)}
train | step: 1800 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.46148437),
'learning_rate': np.float32(0.071288966),
'top_5_accuracy': np.float32(0.90601563),
'training_loss': np.float32(1.7271374)}
train | step: 1900 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.46515626),
'learning_rate': np.float32(0.068406224),
'top_5_accuracy': np.float32(0.9024219),
'training_loss': np.float32(1.7165723)}
train | step: 2000 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.4709375),
'learning_rate': np.float32(0.06545085),
'top_5_accuracy': np.float32(0.9007813),
'training_loss': np.float32(1.7039804)}
eval | step: 2000 | running 78 steps of evaluation...
eval | step: 2000 | steps/sec: 77.9 | eval time: 1.0 sec | output:
{'accuracy': np.float32(0.58683896),
'steps_per_second': 77.9400533170698,
'top_5_accuracy': np.float32(0.9572316),
'validation_loss': np.float32(1.3775176)}
train | step: 2000 | training until step 3000...
train | step: 2100 | steps/sec: 22.3 | output:
{'accuracy': np.float32(0.4884375),
'learning_rate': np.float32(0.062434502),
'top_5_accuracy': np.float32(0.9075781),
'training_loss': np.float32(1.6590232)}
train | step: 2200 | steps/sec: 28.7 | output:
{'accuracy': np.float32(0.49007812),
'learning_rate': np.float32(0.059369065),
'top_5_accuracy': np.float32(0.9102344),
'training_loss': np.float32(1.6366767)}
train | step: 2300 | steps/sec: 28.7 | output:
{'accuracy': np.float32(0.5011719),
'learning_rate': np.float32(0.056266654),
'top_5_accuracy': np.float32(0.9210156),
'training_loss': np.float32(1.6063051)}
train | step: 2400 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.5144531),
'learning_rate': np.float32(0.053139526),
'top_5_accuracy': np.float32(0.9225781),
'training_loss': np.float32(1.5878946)}
train | step: 2500 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.5239844),
'learning_rate': np.float32(0.049999997),
'top_5_accuracy': np.float32(0.9217188),
'training_loss': np.float32(1.5683006)}
train | step: 2600 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.5221094),
'learning_rate': np.float32(0.04686048),
'top_5_accuracy': np.float32(0.9235156),
'training_loss': np.float32(1.5680088)}
train | step: 2700 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.5175),
'learning_rate': np.float32(0.043733336),
'top_5_accuracy': np.float32(0.923125),
'training_loss': np.float32(1.5648788)}
train | step: 2800 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.5253125),
'learning_rate': np.float32(0.040630933),
'top_5_accuracy': np.float32(0.9225781),
'training_loss': np.float32(1.5397162)}
train | step: 2900 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.53375),
'learning_rate': np.float32(0.037565507),
'top_5_accuracy': np.float32(0.925625),
'training_loss': np.float32(1.5342889)}
train | step: 3000 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.53648436),
'learning_rate': np.float32(0.034549143),
'top_5_accuracy': np.float32(0.9266406),
'training_loss': np.float32(1.5213569)}
eval | step: 3000 | running 78 steps of evaluation...
eval | step: 3000 | steps/sec: 77.5 | eval time: 1.0 sec | output:
{'accuracy': np.float32(0.65635014),
'steps_per_second': 77.49399754742826,
'top_5_accuracy': np.float32(0.9708534),
'validation_loss': np.float32(1.1844783)}
train | step: 3000 | training until step 4000...
train | step: 3100 | steps/sec: 22.3 | output:
{'accuracy': np.float32(0.54445314),
'learning_rate': np.float32(0.03159377),
'top_5_accuracy': np.float32(0.9296875),
'training_loss': np.float32(1.5038574)}
train | step: 3200 | steps/sec: 28.9 | output:
{'accuracy': np.float32(0.5450781),
'learning_rate': np.float32(0.028711034),
'top_5_accuracy': np.float32(0.92445314),
'training_loss': np.float32(1.4947772)}
train | step: 3300 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.5601562),
'learning_rate': np.float32(0.025912309),
'top_5_accuracy': np.float32(0.93046874),
'training_loss': np.float32(1.4587433)}
train | step: 3400 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.5546875),
'learning_rate': np.float32(0.023208654),
'top_5_accuracy': np.float32(0.9333594),
'training_loss': np.float32(1.4607594)}
train | step: 3500 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.55445313),
'learning_rate': np.float32(0.020610739),
'top_5_accuracy': np.float32(0.9338281),
'training_loss': np.float32(1.460011)}
train | step: 3600 | steps/sec: 28.7 | output:
{'accuracy': np.float32(0.55789065),
'learning_rate': np.float32(0.018128792),
'top_5_accuracy': np.float32(0.93257815),
'training_loss': np.float32(1.4650867)}
train | step: 3700 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.5684375),
'learning_rate': np.float32(0.015772644),
'top_5_accuracy': np.float32(0.9353125),
'training_loss': np.float32(1.4296772)}
train | step: 3800 | steps/sec: 28.8 | output:
{'accuracy': np.float32(0.57585937),
'learning_rate': np.float32(0.01355157),
'top_5_accuracy': np.float32(0.93859375),
'training_loss': np.float32(1.4083066)}
train | step: 3900 | steps/sec: 29.1 | output:
{'accuracy': np.float32(0.5742969),
'learning_rate': np.float32(0.011474336),
'top_5_accuracy': np.float32(0.935),
'training_loss': np.float32(1.4144776)}
train | step: 4000 | steps/sec: 29.2 | output:
{'accuracy': np.float32(0.57921875),
'learning_rate': np.float32(0.009549147),
'top_5_accuracy': np.float32(0.9351562),
'training_loss': np.float32(1.408786)}
eval | step: 4000 | running 78 steps of evaluation...
eval | step: 4000 | steps/sec: 78.5 | eval time: 1.0 sec | output:
{'accuracy': np.float32(0.7011218),
'steps_per_second': 78.51906061018433,
'top_5_accuracy': np.float32(0.97696316),
'validation_loss': np.float32(1.0532619)}
train | step: 4000 | training until step 5000...
train | step: 4100 | steps/sec: 22.6 | output:
{'accuracy': np.float32(0.573125),
'learning_rate': np.float32(0.0077836006),
'top_5_accuracy': np.float32(0.93953127),
'training_loss': np.float32(1.4169648)}
train | step: 4200 | steps/sec: 29.3 | output:
{'accuracy': np.float32(0.58554685),
'learning_rate': np.float32(0.0061846706),
'top_5_accuracy': np.float32(0.9365625),
'training_loss': np.float32(1.3933432)}
train | step: 4300 | steps/sec: 29.2 | output:
{'accuracy': np.float32(0.58367187),
'learning_rate': np.float32(0.0047586444),
'top_5_accuracy': np.float32(0.93773437),
'training_loss': np.float32(1.3891603)}
train | step: 4400 | steps/sec: 29.2 | output:
{'accuracy': np.float32(0.5865625),
'learning_rate': np.float32(0.0035111725),
'top_5_accuracy': np.float32(0.9385156),
'training_loss': np.float32(1.3999504)}
train | step: 4500 | steps/sec: 29.3 | output:
{'accuracy': np.float32(0.58015627),
'learning_rate': np.float32(0.002447176),
'top_5_accuracy': np.float32(0.9363281),
'training_loss': np.float32(1.3999443)}
train | step: 4600 | steps/sec: 29.2 | output:
{'accuracy': np.float32(0.5789844),
'learning_rate': np.float32(0.0015708387),
'top_5_accuracy': np.float32(0.9355469),
'training_loss': np.float32(1.4043095)}
train | step: 4700 | steps/sec: 29.3 | output:
{'accuracy': np.float32(0.5742969),
'learning_rate': np.float32(0.0008856386),
'top_5_accuracy': np.float32(0.93859375),
'training_loss': np.float32(1.4113929)}
train | step: 4800 | steps/sec: 29.2 | output:
{'accuracy': np.float32(0.5894531),
'learning_rate': np.float32(0.00039426386),
'top_5_accuracy': np.float32(0.93671876),
'training_loss': np.float32(1.3889945)}
train | step: 4900 | steps/sec: 29.2 | output:
{'accuracy': np.float32(0.59046876),
'learning_rate': np.float32(9.866357e-05),
'top_5_accuracy': np.float32(0.93890625),
'training_loss': np.float32(1.3809924)}
train | step: 5000 | steps/sec: 29.2 | output:
{'accuracy': np.float32(0.5985938),
'learning_rate': np.float32(0.0),
'top_5_accuracy': np.float32(0.9459375),
'training_loss': np.float32(1.3506776)}
eval | step: 5000 | running 78 steps of evaluation...
eval | step: 5000 | steps/sec: 78.7 | eval time: 1.0 sec | output:
{'accuracy': np.float32(0.7275641),
'steps_per_second': 78.69840238898342,
'top_5_accuracy': np.float32(0.9786659),
'validation_loss': np.float32(0.9903465)}
saved checkpoint to /tmpfs/tmp/tmp9ai4tfkc/ckpt-5000.
I0000 00:00:1793449535.069865 49795 devices.cc:67] Number of eligible GPUs (core count >= 8, compute capability >= 0.0): 4
I0000 00:00:1793449535.070021 49795 single_machine.cc:376] Starting new session
I0000 00:00:1793449535.100207 49795 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 13638 MB memory: -> device: 0, name: Tesla T4, pci bus id: 0000:00:05.0, compute capability: 7.5
I0000 00:00:1793449535.102068 49795 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 13756 MB memory: -> device: 1, name: Tesla T4, pci bus id: 0000:00:06.0, compute capability: 7.5
I0000 00:00:1793449535.103770 49795 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 13756 MB memory: -> device: 2, name: Tesla T4, pci bus id: 0000:00:07.0, compute capability: 7.5
I0000 00:00:1793449535.105474 49795 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 13756 MB memory: -> device: 3, name: Tesla T4, pci bus id: 0000:00:08.0, compute capability: 7.5
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tensorflow/python/ops/nn_ops.py:5261: tensor_shape_from_node_def_name (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.
Instructions for updating:
This API was designed for TensorFlow v1. See https://www.tensorflow.org/guide/migrate for instructions on how to migrate your code to TensorFlow v2.
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tensorflow/python/ops/nn_ops.py:5261: tensor_shape_from_node_def_name (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.
Instructions for updating:
This API was designed for TensorFlow v1. See https://www.tensorflow.org/guide/migrate for instructions on how to migrate your code to TensorFlow v2.
eval | step: 5000 | running 78 steps of evaluation...
eval | step: 5000 | steps/sec: 69.9 | eval time: 1.1 sec | output:
{'accuracy': np.float32(0.7275641),
'steps_per_second': 69.94087400301522,
'top_5_accuracy': np.float32(0.9786659),
'validation_loss': np.float32(0.9903465)}
```
# tf.keras.utils.plot_model(model, show_shapes=True)
Print the `accuracy`, `top_5_accuracy`, and `validation_loss` evaluation metrics.
for key, value in eval_logs.items():
if isinstance(value, tf.Tensor):
value = value.numpy()
print(f'{key:20}: {value:.3f}')
```
accuracy : 0.728
top_5_accuracy : 0.979
validation_loss : 0.990
steps_per_second : 69.941
```
Run a batch of the processed training data through the model, and view the results
for images, labels in task.build_inputs(exp_config.task.train_data).take(1):
predictions = model.predict(images)
predictions = tf.argmax(predictions, axis=-1)
show_batch(images, labels, tf.cast(predictions, tf.int32))
if device=='CPU':
plt.suptitle('The model was only trained for a few steps, it is not expected to do well.')
```
2026-10-31 12:25:40.425029: W tensorflow/core/kernels/data/cache_dataset_ops.cc:917] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.
4/4 [==============================] - 3s 9ms/step
2026-10-31 12:25:43.954277: W tensorflow/core/kernels/data/cache_dataset_ops.cc:917] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.
```
## Export a SavedModel
The [`keras.Model`](https://www.tensorflow.org/api_docs/python/tf/keras/Model) object returned by [`train_lib.run_experiment`](https://www.tensorflow.org/api_docs/python/tfm/core/train_lib/run_experiment) expects the data to be normalized by the dataset loader using the same mean and variance statiscics in [`preprocess_ops.normalize_image(image, offset=MEAN_RGB, scale=STDDEV_RGB)`](https://www.tensorflow.org/api_docs/python/tfm/vision/preprocess_ops/normalize_image). This export function handles those details, so you can pass [`tf.uint8`](https://www.tensorflow.org/api_docs/python/tf#uint8) images and get the correct results.
# Saving and exporting the trained model
export_saved_model_lib.export_inference_graph(
input_type='image_tensor',
batch_size=1,
input_image_size=[32, 32],
params=exp_config,
checkpoint_path=tf.train.latest_checkpoint(model_dir),
export_dir='./export/')
```
INFO:tensorflow:Assets written to: ./export/assets
INFO:tensorflow:Assets written to: ./export/assets
```
Test the exported model.
# Importing SavedModel
imported = tf.saved_model.load('./export/')
model_fn = imported.signatures['serving_default']
Visualize the predictions.
plt.figure(figsize=(10, 10))
for data in tfds.load('cifar10', split='test').batch(12).take(1):
predictions = []
for image in data['image']:
index = tf.argmax(model_fn(image[tf.newaxis, ...])['logits'], axis=1)[0]
predictions.append(index)
show_batch(data['image'], data['label'], predictions)
if device=='CPU':
plt.suptitle('The model was only trained for a few steps, it is not expected to do better than random.')
```
2026-10-31 12:25:55.816794: W tensorflow/core/kernels/data/cache_dataset_ops.cc:917] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.
```