Supported reference models

Cloud TPU provides a set of reference models that are optimized for fast and accurate training.

Cloud TPU supports the following major and minor framework releases of TensorFlow, PyTorch, and JAX/FLAX. TensorFlow release numbering has changed with release 2.5.0. Major TensorFlow release numbers end with '0' and all patch release numbers end with numbers greater than '0'. For example, TF 2.10.0 is a major release and TF 2.10.1 is a minor release. In order to run the latest-supported framework version, check to see if there are any patch releases to the major release. If so, you can run the latest-supported patch release rather than the major release.

Framework	Major version	Model category	Reference models	Supported versions
TensorFlow	2.x	Image classification	ResNet-2.x, MNIST-2.x, EfficientNet-2.x	2.7.0, 2.7.3, 2.8.0, 2.8.2, 2.8.4, 2.9.1, 2.9.3, 2.10.0, 2.10.1, 2.11.0
		Language modeling	Transformer-2.x, BERT-2.x	2.7.0, 2.7.3, 2.8.0, 2.8.2, 2.8.4, 2.9.1, 2.9.3, 2.10.0, 2.10.1, 2.11.0
		Object detection	RetinaNet-2.x	2.7.0, 2.7.3, 2.8.0, 2.8.2, 2.8.4, 2.9.1, 2.9.3, 2.10.0, 2.10.1, 2.11.0
		Image segmentation	Mask-RCNN-2.x	2.7.0, 2.7.3, 2.8.0, 2.8.2, 2.8.4, 2.9.1, 2.9.3, 2.10.0, 2.10.1, 2.11.0
		Image segmentation	ShapeMask-2.x	2.7.0, 2.7.3, 2.8.0, 2.8.2, 2.8.4, 2.9.1, 2.9.3, 2.10.0, 2.10.1, 2.11.0
		Recommendation systems	DLRM-2.x, DCN-2.x	2.7.0, 2.7.3, 2.8.0, 2.8.2, 2.8.4, 2.9.1, 2.9.3, 2.10.0, 2.10.1, 2.11.0
		Recommendation systems	NCF-2.x	2.7.0, 2.7.3, 2.8.0, 2.8.2, 2.8.4, 2.9.1, 2.9.3, 2.10.0, 2.10.1, 2.11.0
PyTorch	1.x	Image classification	ResNet-PyTorch	1.13
		Language modeling	FairSeq Transformer, FairSeq RoBERTa Wav2Vec2	1.13
		Speech recognition	Wav2Vec2	1.13
		Recommendation systems	DLRM	1.13
JAX/FLAX	latest	Image classification	ResNet50	latest