Running test calculations in DDP mode with multiple GPUs with PyTorchLightning

Question

I have a model which I try to use with trainer in DDP mode.

import pytorch_lightning as pl
import torch
import torchvision
from torchmetrics import Accuracy
class Model(pl.LightningModule):

    def __init__(
        self,
        model_name: str,
        num_classes: int,
        model_hparams: Dict["str", Union[str, int]],
        optimizer_name: str,
        optimizer_hparams: Dict["str", Union[str, int]],
    ):

        super().__init__()
        self.save_hyperparameters()
        self.model = torchvision.resnet18(num_classes=num_classes, **model_hparams)

        self.loss_module = CrossEntropyLoss()
        self.example_input_array = torch.zeros((1, 3, 512, 512), dtype=torch.float32)
        # Trying to use in DDP mode
        self.test_accuracy = Accuracy(num_classes=num_classes)

    def forward(self, imgs) -> Tensor:
        return self.model(imgs)

    # 

    def test_step(self, batch, batch_idx):
        imgs, labels = batch
        preds = self.model(imgs)
        self.test_accuracy.update(preds, labels)
        return labels, preds.argmax(dim=-1)

    def test_epoch_end(self, outputs) -> None:
        num_classes = self.hparams.num_classes
        # Creates table of correct and incorrect predictions
        results = torch.zeros((num_classes, num_classes))
        for output in outputs:
            for label, prediction in zip(*output):
                results[int(label), int(prediction)] += 1
        # Total accuracy. This and `compute` are identical in 1 GPU training
        acc = results.diag().sum() / results.sum()
        self.log("test_acc", self.test_accuracy.compute())
        print(results) # This prints twice

and trainer

trainer = pl.Trainer(
        gpus=torch.cuda.device_count(),
        max_epochs=180,
        callbacks=callbacks,
        strategy="ddp",
        auto_scale_batch_size="binsearch",
    )

However, I get as prints from test

tensor([[0., 0., 0., 0., 0., 5.],
        [0., 7., 0., 0., 0., 0.],
        [0., 3., 0., 0., 0., 2.],
        [0., 3., 0., 0., 0., 0.],
        [0., 3., 0., 0., 0., 2.],
        [0., 1., 0., 0., 0., 4.]])tensor([[0., 0., 0., 0., 0., 6.],
        [0., 2., 0., 0., 0., 0.],
        [0., 4., 0., 0., 0., 2.],
        [0., 2., 0., 0., 0., 1.],
        [0., 3., 0., 0., 0., 2.],
        [0., 5., 0., 0., 0., 3.]])

Also

trainer.fit(model, datamodule=datamodule)
test_results = trainer.test(model, datamodule=datamodule)
print(test_results) 
# [{'test_acc': 0.18333333730697632}]
# [{'test_acc': 0.18333333730697632}]

where I would only expect single tensor to be printed. How can I make my calculations over all test predictions rather than by GPU and return the table I create in test_epoch_end from those predictions? I interpreted the documentation as *_epoch_end being executed only on single GPU and am quite lost.

user_17670380 · Answer

I think you should use following techniques:

test_epoch_end: In ddp mode, every gpu runs same code in this method. So each gpu computes metric on partial batch not whole batches. You need to synchronize metric and collect to rank==0 gpu to compute evaluation metric on entire dataset.
torch.distributed.reduce: This method collects and calculate tensors across distributed gpu devices. (docs)
self.trainer.is_global_zero: This flag will be true for rank==0
What is best way to manually compute metric over test set? you should check docs

Using mentioned techniques, you can compute metric over entire dataset and use results tensor after .test. Here is snippet:

import os

import torch
from torch.utils.data import DataLoader

from torchvision import models, transforms
from torchvision.datasets import CIFAR10

from pytorch_lightning import LightningModule, LightningDataModule, Trainer

os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'


class CIFAR(LightningDataModule):
    def __init__(self, img_size=32, batch_size=32):
        super().__init__()
        self.img_size = img_size if isinstance(img_size, tuple) else (img_size, img_size)
        self.batch_size = batch_size

        self.test_transforms = transforms.Compose([
            transforms.Resize(self.img_size),
            transforms.CenterCrop(self.img_size),
            transforms.ToTensor(),
            transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
        ])

    def prepare_data(self) -> None:
        CIFAR10(root='data', train=True, download=True)
        CIFAR10(root='data', train=False, download=True)
    
    def setup(self, stage=None):
        self.test_ds = CIFAR10(root='data', train=False, download=False, transform=self.test_transforms)

    def test_dataloader(self):
        return DataLoader(self.test_ds, num_workers=4, batch_size=self.batch_size, shuffle=False)


class BasicModule(LightningModule):
    def __init__(self):
        super().__init__()
        self.model = models.resnet18(num_classes=10, pretrained=False)

    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        return y, y_hat.argmax(dim=-1)

    def test_epoch_end(self, outputs):
        results = torch.zeros((10, 10)).to(self.device)
        for output in outputs:
            for label, prediction in zip(*output):
                results[int(label), int(prediction)] += 1
        torch.distributed.reduce(results, 0, torch.distributed.ReduceOp.SUM)
        acc = results.diag().sum() / results.sum()
        if self.trainer.is_global_zero:
            self.log("test_metric", acc, rank_zero_only=True)
            self.trainer.results = results
        
    
if __name__ == '__main__':
    data = CIFAR(batch_size=512)
    model = BasicModule()
    trainer = Trainer(max_epochs=2, gpus='0,1', strategy="ddp", precision=16)
    test_results = trainer.test(model, data)
    if trainer.is_global_zero:
        print(test_results)
        print(trainer.results)

Running test calculations in DDP mode with multiple GPUs with PyTorchLightning

1 Answers