| 33 | |
| 34 | |
| 35 | class TensorAverageMeter: |
| 36 | def __init__(self): |
| 37 | self.tensors = [] |
| 38 | |
| 39 | def add(self, x): |
| 40 | if len(x.shape) == 0: |
| 41 | x = x.unsqueeze(0) |
| 42 | self.tensors.append(x) |
| 43 | |
| 44 | def mean(self): |
| 45 | if len(self.tensors) == 0: |
| 46 | return 0 |
| 47 | cat = torch.cat(self.tensors, dim=0) |
| 48 | if cat.numel() == 0: |
| 49 | return 0 |
| 50 | else: |
| 51 | return cat.mean() |
| 52 | |
| 53 | def clear(self): |
| 54 | self.tensors = [] |
| 55 | |
| 56 | def mean_and_clear(self): |
| 57 | mean = self.mean() |
| 58 | self.clear() |
| 59 | return mean |
| 60 | |
| 61 | |
| 62 | class TensorAverageMeterDict: |