Has anyone successfully used anomalib on a TPU instance on Google Colab?

[WestwardWinds]

I haven't seen much in the docs or in the github about xla support in Anomalib. I know, in theory, it should work because I can start the training process on a TPU instance but right as it starts I end up getting different strange errors. So while I would love to troubleshoot my code, first I would like to see if anyone has been using tpu instances to train models with Anomalib. I will share some examples of lightning working correctly with a TPU instance and then a small number of errors I've encountered trying to make Anomalib work.

Anomalib version: 1.2.0.dev0

I can access the TPU's directly using lightning EX:

import os
import torch
import lightning.pytorch as pl
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

class SimpleModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(10, 1)
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.layer(x)
        loss = nn.functional.mse_loss(y_hat, y)
        self.log('train_loss', loss)
        return loss
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

def main():
    # Create a simple dataset
    train_data = torch.randn(100, 10)
    train_labels = torch.randn(100, 1)
    train_dataset = TensorDataset(train_data, train_labels)

    # Initialize the model and data
    model = SimpleModel()
    train_dataloader = DataLoader(train_dataset, batch_size=32)

    # Configure the trainer
    trainer = pl.Trainer(
        max_epochs=5,
    )

    # Train the model
    trainer.fit(model, train_dataloader)

    print("Training completed successfully on TPU!")

if __name__ == "__main__":
    main()

Outputs:

INFO: GPU available: False, used: False
INFO:lightning.pytorch.utilities.rank_zero:GPU available: False, used: False
INFO: TPU available: True, using: 8 TPU cores
INFO:lightning.pytorch.utilities.rank_zero:TPU available: True, using: 8 TPU cores
INFO: HPU available: False, using: 0 HPUs
INFO:lightning.pytorch.utilities.rank_zero:HPU available: False, using: 0 HPUs
INFO: 
  | Name  | Type   | Params | Mode 
-----------------------------------------
0 | layer | Linear | 11     | train
-----------------------------------------
11        Trainable params
0         Non-trainable params
11        Total params
0.000     Total estimated model params size (MB)
1         Modules in train mode
0         Modules in eval mode
INFO:lightning.pytorch.callbacks.model_summary:
  | Name  | Type   | Params | Mode 
-----------------------------------------
0 | layer | Linear | 11     | train
-----------------------------------------
11        Trainable params
0         Non-trainable params
11        Total params
0.000     Total estimated model params size (MB)
1         Modules in train mode
0         Modules in eval mode
Epoch 4: 100%
 1/1 [00:00<00:00, 25.14it/s, v_num=0]
Training completed successfully on TPU!

But anything I try to do using Anomalib encounters a lot of errors.

EX:

# Import the required modules
from anomalib.data import MVTec
from anomalib.models import Patchcore
from anomalib.engine import Engine

# Initialize the datamodule, model and engine
datamodule = MVTec()
model = Patchcore()
engine = Engine()

# Train the model
engine.fit(datamodule=datamodule, model=model)

Result:

INFO: GPU available: False, used: False
INFO:lightning.pytorch.utilities.rank_zero:GPU available: False, used: False
INFO: TPU available: True, using: 1 TPU cores
INFO:lightning.pytorch.utilities.rank_zero:TPU available: True, using: 1 TPU cores
INFO: HPU available: False, using: 0 HPUs
INFO:lightning.pytorch.utilities.rank_zero:HPU available: False, using: 0 HPUs
WARNING:anomalib.metrics.f1_score:F1Score class exists for backwards compatibility. It will be removed in v1.1. Please use BinaryF1Score from torchmetrics instead
WARNING:anomalib.metrics.f1_score:F1Score class exists for backwards compatibility. It will be removed in v1.1. Please use BinaryF1Score from torchmetrics instead
/usr/local/lib/python3.10/dist-packages/lightning/pytorch/core/optimizer.py:182: `LightningModule.configure_optimizers` returned `None`, this fit will run with no optimizer
INFO: 
  | Name                  | Type                     | Params | Mode 
---------------------------------------------------------------------------
0 | model                 | PatchcoreModel           | 24.9 M | train
1 | _transform            | Compose                  | 0      | train
2 | normalization_metrics | MetricCollection         | 0      | train
3 | image_threshold       | F1AdaptiveThreshold      | 0      | train
4 | pixel_threshold       | F1AdaptiveThreshold      | 0      | train
5 | image_metrics         | AnomalibMetricCollection | 0      | train
6 | pixel_metrics         | AnomalibMetricCollection | 0      | train
---------------------------------------------------------------------------
24.9 M    Trainable params
0         Non-trainable params
24.9 M    Total params
99.450    Total estimated model params size (MB)
17        Modules in train mode
174       Modules in eval mode
INFO:lightning.pytorch.callbacks.model_summary:
  | Name                  | Type                     | Params | Mode 
---------------------------------------------------------------------------
0 | model                 | PatchcoreModel           | 24.9 M | train
1 | _transform            | Compose                  | 0      | train
2 | normalization_metrics | MetricCollection         | 0      | train
3 | image_threshold       | F1AdaptiveThreshold      | 0      | train
4 | pixel_threshold       | F1AdaptiveThreshold      | 0      | train
5 | image_metrics         | AnomalibMetricCollection | 0      | train
6 | pixel_metrics         | AnomalibMetricCollection | 0      | train
---------------------------------------------------------------------------
24.9 M    Trainable params
0         Non-trainable params
24.9 M    Total params
99.450    Total estimated model params size (MB)
17        Modules in train mode
174       Modules in eval mode
Epoch 0:   0%
 0/7 [00:00<?, ?it/s]
---------------------------------------------------------------------------
MisconfigurationException                 Traceback (most recent call last)
[<ipython-input-23-d798772a6e21>](https://localhost:8080/#) in <cell line: 15>()
     13 
     14 # Train the model
---> 15 engine.fit(datamodule=datamodule, model=model)

16 frames
[/usr/local/lib/python3.10/dist-packages/lightning/pytorch/plugins/precision/xla.py](https://localhost:8080/#) in optimizer_step(self, optimizer, model, closure, **kwargs)
     78         # if model.automatic_optimization and skipped_backward:
     79         #     # we lack coverage here so disable this - something to explore if there's demand
---> 80         #     raise MisconfigurationException(
     81         #         "Skipping backward by returning `None` from your `training_step` is not implemented with XLA."
     82         #         " Please, open an issue in `[https://github.com/Lightning-AI/lightning/issues`](https://github.com/Lightning-AI/lightning/issues%60)"

MisconfigurationException: Skipping backward by returning `None` from your `training_step` is not implemented with XLA. Please, open an issue in `[https://github.com/Lightning-AI/lightning/issues`](https://github.com/Lightning-AI/lightning/issues%60) requesting this feature.

Not only is there an error, it is only accessing 1 TPU core instead of the 8 I get using lightning directly.

This is not the only error I've gotten, I wanted to see if I could get a barebones Anomalib example working. In my own projects I've gotten access errors such as this:

model.safetensors: 100%|████████████████████████████████████████████████████████████████████████████████████████████| 46.8M/46.8M [00:00<00:00, 169MB/s]
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1723780056.150404   12260 pjrt_api.cc:100] GetPjrtApi was found for tpu at /usr/local/lib/python3.10/dist-packages/torch_xla/lib/libtpu.so
I0000 00:00:1723780056.150462   12260 pjrt_api.cc:79] PJRT_Api is set for device type tpu
I0000 00:00:1723780056.150467   12260 pjrt_api.cc:146] The PJRT plugin has PJRT API version 0.46. The framework PJRT API version is 0.46.
E0816 03:47:36.264905486   12549 oauth2_credentials.cc:176]            Call to http server ended with error 404 [].
 File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/connectors/accelerator_connector.py", line 153, in __init__
    self._strategy_flag = self._choose_strategy()
  File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/trainer/connectors/accelerator_connector.py", line 436, in _choose_strategy
    return SingleDeviceXLAStrategy(device=self._parallel_devices[0])
  File "/usr/local/lib/python3.10/dist-packages/lightning/pytorch/strategies/single_xla.py", line 53, in __init__
    device=xm.xla_device(device),
  File "/usr/local/lib/python3.10/dist-packages/torch_xla/core/xla_model.py", line 212, in xla_device
    return runtime.xla_device(n, devkind)
  File "/usr/local/lib/python3.10/dist-packages/torch_xla/runtime.py", line 95, in wrapper
    return fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch_xla/runtime.py", line 126, in xla_device
    devices = xm.get_xla_supported_devices(devkind=devkind)
  File "/usr/local/lib/python3.10/dist-packages/torch_xla/core/xla_model.py", line 99, in get_xla_supported_devices
    devices = torch_xla._XLAC._xla_get_devices()
RuntimeError: Bad StatusOr access: UNKNOWN: TPU initialization failed: open(/dev/accel0): Operation not permitted: Operation not permitted; Couldn't open device: /dev/accel0; Unable to create Node RegisterInterface for node 0, config: go/nodeserialize  device_path: "/dev/accel0" mode: KERNEL debug_data_directory: "" dump_anomalies_only: true crash_in_debug_dump: false allow_core_dump: true; could not create driver instance

And some others related to cuda to xla conversion, I think related to the dataset creation but I don't remember.