File size: 1,662 Bytes
d008243 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
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
|