Kaggle TPU

https://www.kaggle.com/docs/tpu
Admin
Kaggle TPUs are hardware accelerators specialized in deep learning tasks, enabling users to train...
Tokens:5,308
Snippets:42
Trust Score:9.8
Update:3 weeks ago
Show doc for...
Tokens:
Raw
### 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)
```
Tokens:
Raw