Compute the average and standard deviation of a dictionary of values.
| 545 | |
| 546 | |
| 547 | class StatisticsCollector: |
| 548 | """ |
| 549 | Compute the average and standard deviation of a dictionary |
| 550 | of values. |
| 551 | """ |
| 552 | |
| 553 | def __init__(self, convert_to_float: bool = True): |
| 554 | """ |
| 555 | Args: |
| 556 | convert_to_float: |
| 557 | whether to convert input values (from Tensor, ndarray) to float |
| 558 | """ |
| 559 | self.convert_to_float = convert_to_float |
| 560 | self.x_sum_dict = dict() |
| 561 | self.x2_sum_dict = dict() |
| 562 | self.x_count_dict = dict() |
| 563 | |
| 564 | def record( |
| 565 | self, |
| 566 | val_dict: T.Dict[str, float], |
| 567 | ): |
| 568 | """ |
| 569 | Record the values in val_dict |
| 570 | |
| 571 | Args: |
| 572 | val_dict: |
| 573 | a dictionary containing the floats to compute statistics. |
| 574 | """ |
| 575 | |
| 576 | for key, val in val_dict.items(): |
| 577 | if isinstance(val, torch.Tensor) and val.numel() > 1: |
| 578 | continue |
| 579 | |
| 580 | if self.convert_to_float: |
| 581 | if isinstance(val, torch.Tensor): |
| 582 | val = val.detach().cpu().item() |
| 583 | elif isinstance(val, np.ndarray): |
| 584 | val = val.item() |
| 585 | elif isinstance(val, int): |
| 586 | val = float(val) |
| 587 | elif isinstance(val, float): |
| 588 | pass |
| 589 | else: |
| 590 | raise NotImplementedError(f"{type(val)}") |
| 591 | |
| 592 | if key not in self.x_sum_dict: |
| 593 | self.x_sum_dict[key] = val |
| 594 | self.x2_sum_dict[key] = val ** 2 |
| 595 | self.x_count_dict[key] = 1 |
| 596 | else: |
| 597 | self.x_sum_dict[key] = self.x_sum_dict[key] + val |
| 598 | self.x2_sum_dict[key] = self.x2_sum_dict[key] + val ** 2 |
| 599 | self.x_count_dict[key] = self.x_count_dict[key] + 1 |
| 600 | |
| 601 | def compute_statistics(self) -> T.Dict[str, T.Dict[str, float]]: |
| 602 | """Compute the statistics. |
| 603 | |
| 604 | Returns: |