| 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 |
|
|
