|
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): |
|
|
|
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 |
|
|
