| | from typing import List, Optional, Union |
| |
|
| | import torch |
| | import torch.distributed |
| | from torchpack import distributed |
| |
|
| | from utils.misc import list_mean, list_sum |
| |
|
| | __all__ = ["ddp_reduce_tensor", "DistributedMetric"] |
| |
|
| |
|
| | def ddp_reduce_tensor( |
| | tensor: torch.Tensor, reduce="mean" |
| | ) -> Union[torch.Tensor, List[torch.Tensor]]: |
| | tensor_list = [torch.empty_like(tensor) for _ in range(distributed.size())] |
| | torch.distributed.all_gather(tensor_list, tensor.contiguous(), async_op=False) |
| | if reduce == "mean": |
| | return list_mean(tensor_list) |
| | elif reduce == "sum": |
| | return list_sum(tensor_list) |
| | elif reduce == "cat": |
| | return torch.cat(tensor_list, dim=0) |
| | elif reduce == "root": |
| | return tensor_list[0] |
| | else: |
| | return tensor_list |
| |
|
| |
|
| | class DistributedMetric(object): |
| | """Average metrics for distributed training.""" |
| |
|
| | def __init__(self, name: Optional[str] = None, backend="ddp"): |
| | self.name = name |
| | self.sum = 0 |
| | self.count = 0 |
| | self.backend = backend |
| |
|
| | def update(self, val: Union[torch.Tensor, int, float], delta_n=1): |
| | val *= delta_n |
| | if type(val) in [int, float]: |
| | val = torch.Tensor(1).fill_(val).cuda() |
| | if self.backend == "ddp": |
| | self.count += ddp_reduce_tensor( |
| | torch.Tensor(1).fill_(delta_n).cuda(), reduce="sum" |
| | ) |
| | self.sum += ddp_reduce_tensor(val.detach(), reduce="sum") |
| | else: |
| | raise NotImplementedError |
| |
|
| | @property |
| | def avg(self): |
| | if self.count == 0: |
| | return torch.Tensor(1).fill_(-1) |
| | else: |
| | return self.sum / self.count |
| |
|
