Metrics reduction on distributed TPU setting

[vfdev-5]

🚀 Feature

Ignite will support distributed training on TPU (e.g. #960). Currently, metric's computation is impacted in the same way as for DDP on GPUs.

Idea is to improve metric's computation and reduce internal values as it is done for DDP:

ignite/ignite/metrics/metric.py

Line 92 in 0b36397

def _sync_all_reduce(self, tensor: Union[torch.Tensor, numbers.Number]) -> Union[torch.Tensor, numbers.Number]:

To check if we are running in distributed TPU, we can opt to

# global definition
try:
     import torch_xla.core.xla_model as xm
     on_xla_device = True
except ImportError:
     on_xla_device = False

and if we need to reduce:

xm.xrt_world_size() > 1

This issues depends on #963

[vfdev-5]

The following monkeypatching works to reduce metrics

# We need to monkeypatch base Ignite metric class to work in distributed TPU
# Until full support of TPU on PyTorch-Ignite side : 
# https://github.com/pytorch/ignite/issues/965
import numbers
import torch
import torch_xla.core.xla_model as xm


def _tpu_sync_all_reduce(self, tensor):
    tensor_to_number = False
    if isinstance(tensor, numbers.Number):
        tensor = torch.tensor(tensor, device=self._device, dtype=torch.float)
        tensor_to_number = True

    if isinstance(tensor, torch.Tensor):
        # check if the tensor is at specified device
        if tensor.device != self._device:
            tensor = tensor.to(self._device)
    else:
        raise TypeError("Unhandled input type {}".format(type(tensor)))

    # synchronize and reduce
    xm.all_reduce("sum", [tensor, ])

    if tensor_to_number:
        return tensor.item()
    return tensor


from ignite.metrics import Metric


Metric._sync_all_reduce = _tpu_sync_all_reduce

https://colab.research.google.com/drive/1Gy8bblDyXYBqMI7PuLAcIDKzRZGSoeYl

[vfdev-5]

Follow up about dtype support : pytorch/xla#1952