128 lines
4.0 KiB
Python
128 lines
4.0 KiB
Python
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
|