Trainer.predict returns None when tpu_cores=8

[KijitoraButi]

🐛 Bug

The Trainer.predict error was fixed in this issue, but there was another problem.
When tpuc_cores=8, Trainer.predict returns None.

Please reproduce using the BoringModel

https://colab.research.google.com/drive/1evjVtSbQpVfiWpTsiZknT6PO6Fa7u7sY?usp=sharing

To Reproduce

If you set tpu_cores=8 and run Trainer.predict, None will be returned.

Expected behavior

Predicted values are returned from Trainer.predict.

Environment

• PyTorch Version (e.g., 1.0): 1.8.1+cu101
• OS (e.g., Linux): Linux
• Python version: 3.7.10
• Any other relevant information: Colab notebook with TPU

[justusschock]

cc @kaushikb11 😎

[kaushikb11]

@KijitoraButi return_predictions is not supported by plugins that spawn processes. TPUSpawn Plugin is a subclass of DDPSpawn plugin and spawn processes, it is not supported.

[tchaton]

Dear @KijitoraButi,

Supporting trainer.predict for TPU could result into OOM and we decided it would be unwise.

However, we have a built-in Callback called BasePredictionWriter: https://pytorch-lightning.readthedocs.io/en/stable/extensions/generated/pytorch_lightning.callbacks.BasePredictionWriter.html#pytorch_lightning.callbacks.BasePredictionWriter.

This is meant to enable you to save your predictions as you want.

import torch
from pytorch_lightning.callbacks import BasePredictionWriter

class CustomWriter(BasePredictionWriter):

    def __init__(self, output_dir: str, write_interval: str):
        super().__init__(write_interval)
        self.output_dir

    def write_on_batch_end(
        self, trainer, pl_module: 'LightningModule', prediction: Any, batch_indices: List[int], batch: Any,
        batch_idx: int, dataloader_idx: int
    ):
        # save on batch

    def write_on_epoch_end(
        self, trainer, pl_module: 'LightningModule', predictions: List[Any], batch_indices: List[Any]
    ):
       # save on epoch

Best,
T.C

[stekiri]

Thanks to the help of @ananthsub, I made the prediction work on TPU (see here). @tchaton, I would love to learn more about the design decision to prevent returning predictions for TPUs. Do the OOMs only occur for very large datasets? If so, shouldn't it be up to the user to decide whether his dataset can fit into memory?

I only ask because it could be an awesome addition to the library if the hardware-specific complexity that comes with using a CustomWriter callback is abstracted away for the user. In TensorFlow, this is handled by tf.keras.Model.predict() that returns a numpy matrix containing all predictions.

Maybe there are technical limitations that I am not aware of, I can just speak from the user perspective. So, if this is actually feasible, I'd love to see this being supported.

[stekiri]

@tchaton, thanks for your suggestion to use the BasePredictionWriter. I am struggling to make this work when using Colab with a TPU runtime and tpu_cores=8. What I experience is that only an eighth of the predictions are actually saved. Maybe I am doing something wrong or is this simply not supported?

Here's the code to reproduce the issue, it is based on the TPU tutorial with a few modifications:

Package installation:

!pip install torch==1.9.1 torchtext==0.10.1 torchvision==0.10.1 pytorch-lightning==1.5.8 cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl

Code:

import os
from typing import Any, List, Optional

import torch
import torch.nn.functional as F
from pytorch_lightning import LightningDataModule, LightningModule, Trainer
from pytorch_lightning.callbacks import BasePredictionWriter
from torch import nn
from torch.utils.data import DataLoader, random_split
from torchmetrics.functional import accuracy
from torchvision import transforms
from torchvision.datasets import MNIST

BATCH_SIZE = 1024


class MNISTDataModule(LightningDataModule):
    def __init__(self, data_dir: str = "./"):
        super().__init__()
        self.data_dir = data_dir
        self.transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

        # self.dims is returned when you call dm.size()
        # Setting default dims here because we know them.
        # Could optionally be assigned dynamically in dm.setup()
        self.dims = (1, 28, 28)
        self.num_classes = 10

    def prepare_data(self):
        # download
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage=None):
        mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)
        self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])

        self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=BATCH_SIZE)

    def val_dataloader(self):
        return DataLoader(self.mnist_val, batch_size=BATCH_SIZE)

    def test_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=BATCH_SIZE)

    def predict_dataloader(self):
        return DataLoader(self.mnist_test, batch_size=BATCH_SIZE)


class LitModel(LightningModule):
    def __init__(self, channels, width, height, num_classes, hidden_size=64, learning_rate=2e-4):
        super().__init__()

        self.save_hyperparameters()

        self.model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(channels * width * height, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_size, num_classes),
        )

    def forward(self, x):
        x = self.model(x)
        return F.log_softmax(x, dim=1)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.nll_loss(logits, y)
        preds = torch.argmax(logits, dim=1)
        acc = accuracy(preds, y)
        self.log("val_loss", loss, prog_bar=True)
        self.log("val_acc", acc, prog_bar=True)
        return loss

    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: Optional[int] = None) -> Any:
        x, y = batch
        return self(x)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
        return optimizer


class CustomWriter(BasePredictionWriter):
    def __init__(self, output_dir: str, write_interval: str):
        super().__init__(write_interval)
        self.output_dir = output_dir

    def write_on_batch_end(
            self, trainer, pl_module: LightningModule, prediction: Any, batch_indices: List[int], batch: Any,
            batch_idx: int, dataloader_idx: int):
        torch.save(prediction, os.path.join(self.output_dir, f"{dataloader_idx}_{batch_idx}.pt"))

    def write_on_epoch_end(
            self, trainer, pl_module: LightningModule, predictions: List[Any], batch_indices: List[Any]):
        torch.save(predictions, os.path.join(self.output_dir, "predictions.pt"))


tmp_dir = "/tmp"
dm = MNISTDataModule()
model = LitModel(*dm.size(), dm.num_classes)
prediction_writer = CustomWriter(
    output_dir=tmp_dir,
    write_interval="epoch")
trainer = Trainer(
    tpu_cores=8,
    callbacks=[prediction_writer])

trainer.predict(model=model, datamodule=dm)

written_predictions = torch.load(os.path.join(tmp_dir, 'predictions.pt'))
nb_predictions = sum([t.shape[0] for t in written_predictions[0]])

assert nb_predictions == 10_000

It works correctly using a single core with tpu_cores=[1], but it would definitely be nice if all the cores could be used for the prediction.

[kaushikb11]

@stekiri Could you create an issue for it? So that I could look into it.

[stekiri]

@kaushikb11, created this new issue #11417 to follow-up. Thanks!

[ananthsub]

@stekiri - in your callback, each TPU core is overwriting the same files, the batch/dataloader idx and predictions.pt . So when you're loading them afterwards, you're seeing only a portion of the total (whatever was saved last). Either open the files in append mode, or write to a different file per core and group them together afterwards.

If you open the file in append mode, be sure to close it afterwards at the end of prediction.

if you're partitioning the files, you can use trainer.global_rank

[shashankg7]

Hi,

Can I call CustomWriter only with TPUs? I want to write a custom prediction writer on GPU. I tried adding this callback, but it doesn't work.