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