### Install Torch-XLA for TPU training Source: https://www.kaggle.com/docs/tpu Install Torch-XLA and set up the environment for TPU training in PyTorch. The commands use the nightly version and install required apt packages. After installation, you can run PyTorch code on TPU v3-8, which has 8 cores; to fully utilize the TPU, run code on all 8 cores simultaneously. ```bash # Step 1: Install Torch-XLA (PyTorch with Accelerated Linear Algebra (XLA) support) !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py !python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev # Step 2: Run your PyTorch code TPUs (TPU v3-8) have 8 cores, and each core is itself an XLA device. You can run code on a single XLA device, but to take full advantage of the TPU you will want to run your code on all 8 cores simultaneously. For examples that demonstrate how to do this, you can refer to The Ultimate PyTorch TPU Tutorial, I Like Clean TPU Training Kernels and I Can Not Lie, Super Duper Fast PyTorch TPU Kernel, and XLM Roberta Large Pytorch TPU ``` -------------------------------- ### Set up Google Cloud credentials for private Kaggle Datasets Source: https://www.kaggle.com/docs/tpu Initializes Google Cloud credentials for using private Kaggle Datasets with TPUs. Requires enabling 'Google Cloud SDK' in the notebook's Add-ons menu. For public datasets, only Step 3 is needed. ```python # Step 1: Get the credential from the Cloud SDK from kaggle_secrets import UserSecretsClient user_secrets = UserSecretsClient() user_credential = user_secrets.get_gcloud_credential() # Step 2: Set the credentials user_secrets.set_tensorflow_credential(user_credential) # Step 3: Use a familiar call to get the GCS path of the dataset from kaggle_datasets import KaggleDatasets GCS_DS_PATH = KaggleDatasets().get_gcs_path() ``` -------------------------------- ### PyTorch TPU training checklist Source: https://www.kaggle.com/docs/tpu Key steps for using TPUs with PyTorch, covering distributed training with xmp.spawn, model wrapping with MpModelWrapper, sending data to the XLA device, and using ParallelLoader for training. Also includes memory-optimized and standard model save/restore methods, and notes on loading out-of-memory datasets from localhost or GCS in TFRecord format. ```python #1: Startup Script https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py #2: Distributed training function mp_fn xmp.spawn(_mp_fn, nprocs=8, start_method='fork') #3: Instantiate model outside of mp_fn and use MpModelWrapper MX = JigsawModel() => MX = xmp.MpModelWrapper(JigsawModel()) #4: Send model to TPU device device = xm.xla_device() model = MX.to(device) #5: Changes to training loop: send data to device ids = ids.to(device, dtype=torch.long) token_type_ids = token_type_ids.to(device, dtype=torch.long) mask = mask.to(device, dtype=torch.long) targets = targets.to(device, dtype=torch.float) #6: Printing messages xm.master_print #7: Loading data train_dataset = … # user-defined, can be outside of mp_fn # in mp_fn: train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, num_replicas=xm.xrt_world_size(),rank=xm.get_ordinal(), …) train_data_loader = torch.utils.data.DataLoader(train_dataset, sampler=train_sampler, …) #8: Training on data for epoch in range(EPOCHS): para_loader = pl.ParallelLoader(train_data_loader, [device]) train_fn(para_loader.per_device_loader(device), …) #9: Results from TPU xm.mesh_reduce #10: Model save / restore (memory-optimized) import torch_xla.utils.serialization as xser xser.save(model.state_dict(), f"model.bin", master_only=True) model.load_state_dict(xser.load(f"model.bin")) #11: Model save / restore (PyTorch standard) torch_xla.core.xla_model.save torch.load(...) #12: Out of memory datasets: Can be loaded from localhost Of loaded from GCS in TFRecord format, a TFRecords PyTorch loader exists ``` -------------------------------- ### Load TPU model from local disk with tf.saved_model.LoadOptions Source: https://www.kaggle.com/docs/tpu Use this snippet to load a TPU model from local disk. The model must be loaded within the strategy scope, and the experimental_io_device option is required for local loading. ```python with strategy.scope(): load_locally = tf.saved_model.LoadOptions(experimental_io_device='/job:localhost') model = tf.keras.models.load_model('./model', options=load_locally) # loading in Tensorflow's "SavedModel" format ``` -------------------------------- ### Load TPU-trained model for GPU inference Source: https://www.kaggle.com/docs/tpu Step 3 of the TPU workaround: load the saved model into an inference notebook and run inference using a GPU. This notebook can be submitted to competitions because it only uses a GPU. ```python # You can now load your model and run inference using a GPU in this notebook. # Because this notebook only uses a GPU, you can submit it to competitions model = tf.keras.models.load_model('../input/yourDataset/model.h5') ``` -------------------------------- ### Scale batch size with TPU core count using num_replicas_in_sync Source: https://www.kaggle.com/docs/tpu Scale the batch size with the number of TPU cores. With a TPUStrategy on a single TPU v3-8 the core count is 8, so this yields a batch size of 128 per core, the recommended rule of thumb for keeping the 128x128 hardware matrix multipliers busy. ```python BATCH_SIZE = 16 * tpu_strategy.num_replicas_in_sync ``` -------------------------------- ### Connect to TPU and instantiate TPUStrategy Source: https://www.kaggle.com/docs/tpu This boilerplate connects to a TPU and instantiates a distribution strategy. The strategy scope instructs TensorFlow to instantiate all model variables in TPU memory, and the TPUClusterResolver.connect() call automatically enters the TPU device scope. ```python # connect to a TPU and instantiate a distribution strategy tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='local') tf.tpu.experimental.initialize_tpu_system(tpu) tpu_strategy = tf.distribute.TPUStrategy(tpu) # instantiate the model in the strategy scope with tpu_strategy.scope(): model = tf.keras.Sequential( … ) ``` -------------------------------- ### Load TensorFlow Hub model to TPU with load_options Source: https://www.kaggle.com/docs/tpu Use this snippet to load a model from TensorFlow Hub directly to TPU. The load_options parameter with experimental_io_device is required for local loading. ```python import tensorflow_hub as hub with strategy.scope(): load_locally = tf.saved_model.LoadOptions(experimental_io_device='/job:localhost') pretrained_model = hub.KerasLayer('https://tfhub.dev/tensorflow/efficientnet/b6/feature-vector/1', trainable=True, input_shape=[512,512,3], load_options=load_locally) ``` -------------------------------- ### Enable TPU training in TensorFlow Keras with TPUStrategy Source: https://www.kaggle.com/docs/tpu Enable TPU training in TensorFlow Keras after switching the notebook accelerator to TPU v3-8. Use TPUClusterResolver to locate the network-connected TPU, initialize the system, create a TPUStrategy, and instantiate the model inside the strategy scope so it is created on the TPU. Model size is limited by TPU RAM, not VM memory. ```python # detect and init the TPU tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # instantiate a distribution strategy tf.tpu.experimental.initialize_tpu_system(tpu) tpu_strategy = tf.distribute.TPUStrategy(tpu) # instantiating the model in the strategy scope creates the model on the TPU with tpu_strategy.scope(): model = tf.keras.Sequential( … ) # define your model normally model.compile( … ) # train model normally model.fit(training_dataset, epochs=EPOCHS, steps_per_epoch=…) ``` -------------------------------- ### Save TPU model locally with tf.saved_model.SaveOptions Source: https://www.kaggle.com/docs/tpu Use this snippet to save a TPU model to local disk. The experimental_io_device option is required for local saving; it can be omitted when writing to GCS because TPUs have direct access to GCS. This option does nothing on GPUs. ```python save_locally = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost') model.save('./model', options=save_locally) # saving in Tensorflow's "SavedModel" format ``` -------------------------------- ### Write checkpoints locally from TPU model with ModelCheckpoint Source: https://www.kaggle.com/docs/tpu Use this snippet to write checkpoints locally from a TPU model. The experimental_io_device option is required for local checkpoint saving. ```python save_locally = tf.saved_model.SaveOptions(experimental_io_device='/job:localhost') checkpoints_cb = tf.keras.callbacks.ModelCheckpoint('./checkpoints', options=save_locally) model.fit(…, callbacks=[checkpoints_cb]) ``` -------------------------------- ### Parallel TFRecord streaming with num_parallel_reads and experimental_deterministic Source: https://www.kaggle.com/docs/tpu Enables parallel streaming from multiple TFRecord files by setting num_parallel_reads=AUTO and disabling deterministic order with experimental_deterministic=False. This improves throughput on TPUs by reading from multiple files concurrently. ```python AUTO = tf.data.experimental.AUTOTUNE ignore_order = tf.data.Options() ignore_order.experimental_deterministic = False # On Kaggle you can also use KaggleDatasets().get_gcs_path() to obtain the GCS path of a Kaggle dataset filenames = tf.io.gfile.glob("gs://flowers-public/tfrecords-jpeg-512x512/*.tfrec") # list files on GCS dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO) dataset = dataset.with_options(ignore_order) dataset = dataset.map(...) # TFRecord decoding here... ``` -------------------------------- ### Load TFRecord dataset from GCS with tf.data.TFRecordDataset Source: https://www.kaggle.com/docs/tpu Loads a TFRecord dataset from GCS using tf.data.TFRecordDataset. The glob pattern lists all .tfrec files in the specified GCS directory. The map call is a placeholder for TFRecord decoding logic. ```python # On Kaggle you can also use KaggleDatasets().get_gcs_path() to obtain the GCS path of a Kaggle dataset filenames = tf.io.gfile.glob("gs://flowers-public/tfrecords-jpeg-512x512/*.tfrec") # list files on GCS dataset = tf.data.TFRecordDataset(filenames) dataset = dataset.map(...) # TFRecord decoding here... ``` -------------------------------- ### Save TPU-trained model to disk in .h5 format Source: https://www.kaggle.com/docs/tpu Step 1 of the TPU workaround for code competitions: save the trained model to disk in .h5 format. This step will be replaced with an alternative call in TensorFlow 2.3. ```python # Save your model to disk using the .save() functionality. Here we save in .h5 format # This step will be replaced with an alternative call to save models in Tensorflow 2.3 model.save('model.h5') ``` -------------------------------- ### Set steps_per_execution=32 in model.compile() for TPU Source: https://www.kaggle.com/docs/tpu Use the steps_per_execution parameter (available from TensorFlow 2.4) in model.compile() to send multiple batches to the TPU at once. This lowers communication overhead and lets the XLA compiler optimize hardware utilization across batches, so very high batch sizes are no longer required for good performance. ```python model.compile( … , steps_per_execution=32) ``` === COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.