Spaces:
Runtime error
Runtime error
| from enum import Enum | |
| import numpy as np | |
| import torch | |
| import torch.distributed as dist | |
| IGNORE_INDEX = -100 | |
| class Summary(Enum): | |
| NONE = 0 | |
| AVERAGE = 1 | |
| SUM = 2 | |
| COUNT = 3 | |
| class AverageMeter(object): | |
| """Computes and stores the average and current value""" | |
| def __init__(self, name, fmt=":f", summary_type=Summary.AVERAGE): | |
| self.name = name | |
| self.fmt = fmt | |
| self.summary_type = summary_type | |
| self.reset() | |
| def reset(self): | |
| self.val = 0 | |
| self.avg = 0 | |
| self.sum = 0 | |
| self.count = 0 | |
| def update(self, val, n=1): | |
| self.val = val | |
| self.sum += val * n | |
| self.count += n | |
| self.avg = self.sum / self.count | |
| def all_reduce(self): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| if isinstance(self.sum, np.ndarray): | |
| total = torch.tensor( | |
| self.sum.tolist() | |
| + [ | |
| self.count, | |
| ], | |
| dtype=torch.float32, | |
| device=device, | |
| ) | |
| else: | |
| total = torch.tensor( | |
| [self.sum, self.count], dtype=torch.float32, device=device | |
| ) | |
| dist.all_reduce(total, dist.ReduceOp.SUM, async_op=False) | |
| if total.shape[0] > 2: | |
| self.sum, self.count = total[:-1].cpu().numpy(), total[-1].cpu().item() | |
| else: | |
| self.sum, self.count = total.tolist() | |
| self.avg = self.sum / (self.count + 1e-5) | |
| def __str__(self): | |
| fmtstr = "{name} {val" + self.fmt + "} ({avg" + self.fmt + "})" | |
| return fmtstr.format(**self.__dict__) | |
| def summary(self): | |
| fmtstr = "" | |
| if self.summary_type is Summary.NONE: | |
| fmtstr = "" | |
| elif self.summary_type is Summary.AVERAGE: | |
| fmtstr = "{name} {avg:.3f}" | |
| elif self.summary_type is Summary.SUM: | |
| fmtstr = "{name} {sum:.3f}" | |
| elif self.summary_type is Summary.COUNT: | |
| fmtstr = "{name} {count:.3f}" | |
| else: | |
| raise ValueError("invalid summary type %r" % self.summary_type) | |
| return fmtstr.format(**self.__dict__) | |
| def intersectionAndUnionGPU(output, target, K, ignore_index=255): | |
| # 'K' classes, output and target sizes are N or N * L or N * H * W, each value in range 0 to K - 1. | |
| assert output.dim() in [1, 2, 3] | |
| assert output.shape == target.shape | |
| output = output.view(-1) | |
| target = target.view(-1) | |
| output[target == ignore_index] = ignore_index | |
| intersection = output[output == target] | |
| area_intersection = torch.histc(intersection, bins=K, min=0, max=K - 1) | |
| area_output = torch.histc(output, bins=K, min=0, max=K - 1) | |
| area_target = torch.histc(target, bins=K, min=0, max=K - 1) | |
| area_union = area_output + area_target - area_intersection | |
| return area_intersection, area_union, area_target | |
| class ProgressMeter(object): | |
| def __init__(self, num_batches, meters, prefix=""): | |
| self.batch_fmtstr = self._get_batch_fmtstr(num_batches) | |
| self.meters = meters | |
| self.prefix = prefix | |
| def display(self, batch): | |
| entries = [self.prefix + self.batch_fmtstr.format(batch)] | |
| entries += [str(meter) for meter in self.meters] | |
| print("\t".join(entries)) | |
| def display_summary(self): | |
| entries = [" *"] | |
| entries += [meter.summary() for meter in self.meters] | |
| print(" ".join(entries)) | |
| def _get_batch_fmtstr(self, num_batches): | |
| num_digits = len(str(num_batches // 1)) | |
| fmt = "{:" + str(num_digits) + "d}" | |
| return "[" + fmt + "/" + fmt.format(num_batches) + "]" | |
| def dict_to_cuda(input_dict): | |
| for k, v in input_dict.items(): | |
| if isinstance(input_dict[k], torch.Tensor): | |
| input_dict[k] = v.cuda(non_blocking=True) | |
| elif ( | |
| isinstance(input_dict[k], list) | |
| and len(input_dict[k]) > 0 | |
| and isinstance(input_dict[k][0], torch.Tensor) | |
| ): | |
| input_dict[k] = [ele.cuda(non_blocking=True) for ele in v] | |
| return input_dict | |