Computes and stores the average and current value
| 438 | COUNT = 3 |
| 439 | |
| 440 | class AverageMeter(object): |
| 441 | """Computes and stores the average and current value""" |
| 442 | def __init__(self, name, use_accel, fmt=':f', summary_type=Summary.AVERAGE): |
| 443 | self.name = name |
| 444 | self.use_accel = use_accel |
| 445 | self.fmt = fmt |
| 446 | self.summary_type = summary_type |
| 447 | self.reset() |
| 448 | |
| 449 | def reset(self): |
| 450 | self.val = 0 |
| 451 | self.avg = 0 |
| 452 | self.sum = 0 |
| 453 | self.count = 0 |
| 454 | |
| 455 | def update(self, val, n=1): |
| 456 | self.val = val |
| 457 | self.sum += val * n |
| 458 | self.count += n |
| 459 | self.avg = self.sum / self.count |
| 460 | |
| 461 | def all_reduce(self): |
| 462 | if self.use_accel: |
| 463 | device = torch.accelerator.current_accelerator() |
| 464 | else: |
| 465 | device = torch.device("cpu") |
| 466 | total = torch.tensor([self.sum, self.count], dtype=torch.float32, device=device) |
| 467 | dist.all_reduce(total, dist.ReduceOp.SUM, async_op=False) |
| 468 | self.sum, self.count = total.tolist() |
| 469 | self.avg = self.sum / self.count |
| 470 | |
| 471 | def __str__(self): |
| 472 | fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})' |
| 473 | return fmtstr.format(**self.__dict__) |
| 474 | |
| 475 | def summary(self): |
| 476 | fmtstr = '' |
| 477 | if self.summary_type is Summary.NONE: |
| 478 | fmtstr = '' |
| 479 | elif self.summary_type is Summary.AVERAGE: |
| 480 | fmtstr = '{name} {avg:.3f}' |
| 481 | elif self.summary_type is Summary.SUM: |
| 482 | fmtstr = '{name} {sum:.3f}' |
| 483 | elif self.summary_type is Summary.COUNT: |
| 484 | fmtstr = '{name} {count:.3f}' |
| 485 | else: |
| 486 | raise ValueError('invalid summary type %r' % self.summary_type) |
| 487 | |
| 488 | return fmtstr.format(**self.__dict__) |
| 489 | |
| 490 | |
| 491 | class ProgressMeter(object): |