elun15/image-regression

1from enum import Enum2 3import torch4import torch.distributed as dist5 6 7class Summary(Enum):8    NONE = 09    AVERAGE = 110    SUM = 211    COUNT = 312 13 14class AverageMeter(object):15    """Computes and stores the average and current value"""16 17    def __init__(self, name, fmt=":f", summary_type=Summary.AVERAGE):18        self.name = name19        self.fmt = fmt20        self.summary_type = summary_type21        self.reset()22 23    def reset(self):24        self.val = 025        self.avg = 026        self.sum = 027        self.count = 028 29    def update(self, val, n=1):30        self.val = val31        self.sum += val * n32        self.count += n33        self.avg = self.sum / self.count34 35    def all_reduce(self):36        if torch.cuda.is_available():37            device = torch.device("cuda")38        elif torch.backends.mps.is_available():39            device = torch.device("mps")40        else:41            device = torch.device("cpu")42        total = torch.tensor([self.sum, self.count], dtype=torch.float32, device=device)43        dist.all_reduce(total, dist.ReduceOp.SUM, async_op=False)44        self.sum, self.count = total.tolist()45        self.avg = self.sum / self.count46 47    def __str__(self):48        fmtstr = "{name} {val" + self.fmt + "} ({avg" + self.fmt + "})"49        return fmtstr.format(**self.__dict__)50 51    def summary(self):52        fmtstr = ""53        if self.summary_type is Summary.NONE:54            fmtstr = ""55        elif self.summary_type is Summary.AVERAGE:56            fmtstr = "{name} {avg:.3f}"57        elif self.summary_type is Summary.SUM:58            fmtstr = "{name} {sum:.3f}"59        elif self.summary_type is Summary.COUNT:60            fmtstr = "{name} {count:.3f}"61        else:62            raise ValueError("invalid summary type %r" % self.summary_type)63 64        return fmtstr.format(**self.__dict__)65