|
- from collections import defaultdict, deque
- import datetime
- import time
- import torch
- import torch.distributed as dist
- import numpy as np
- import errno
- import os
-
- class SmoothedValue(object):
- """Track a series of values and provide access to smoothed values over a
- window or the global series average.
- """
-
- def __init__(self, window_size=20, fmt=None):
- if fmt is None:
- fmt = "{value:.4f} ({global_avg:.4f})"
- self.deque = deque(maxlen=window_size)
- self.total = 0.0
- self.count = 0
- self.fmt = fmt
-
- def update(self, value, n=1):
- self.deque.append(value)
- self.count += n
- self.total += value * n
-
- def synchronize_between_processes(self):
- """
- Warning: does not synchronize the deque!
- """
- if not is_dist_avail_and_initialized():
- return
- t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
- dist.barrier()
- dist.all_reduce(t)
- t = t.tolist()
- self.count = int(t[0])
- self.total = t[1]
-
- @property
- def median(self):
- d = torch.tensor(list(self.deque))
- return d.median().item()
-
- @property
- def avg(self):
- d = torch.tensor(list(self.deque), dtype=torch.float32)
- return d.mean().item()
-
- @property
- def global_avg(self):
- return self.total / self.count
-
- @property
- def max(self):
- return max(self.deque)
-
- @property
- def value(self):
- return self.deque[-1]
-
- def __str__(self):
- return self.fmt.format(
- median=self.median,
- avg=self.avg,
- global_avg=self.global_avg,
- max=self.max,
- value=self.value)
-
-
- class ConfusionMatrix(object):
- def __init__(self, num_classes):
- self.num_classes = num_classes
- self.mat = None
-
- def update(self, a, b):
- n = self.num_classes
- if self.mat is None:
- # 创建混淆矩阵
- self.mat = torch.zeros((n, n), dtype=torch.int64, device=a.device)
- with torch.no_grad():
- # 寻找GT中为目标的像素索引
-
- k = (a >= 0) & (a < n)
- # 统计像素真实类别a[k]被预测成类别b[k]的个数(这里的做法很巧妙)
- inds = n * a[k].to(torch.int64) + b[k]
- self.mat += torch.bincount(inds, minlength=n ** 2).reshape(n, n)
-
- def reset(self):
- if self.mat is not None:
- self.mat.zero_()
-
- def compute(self):
- h = self.mat.float()
- hist = h.cpu().numpy()
- # 计算全局预测准确率(混淆矩阵的对角线为预测正确的个数)
- acc_global = torch.diag(h).sum() / h.sum()
- # 计算每个类别的准确率
- acc = torch.diag(h) / h.sum(1)
- # 计算每个类别预测与真实目标的iou
- iu = torch.diag(h) / (h.sum(1) + h.sum(0) - torch.diag(h))
-
- # # 计算每个类别的召回率
- # recall = torch.diag(h) / torch.maximum(h.sum(1), torch.tensor(1))
- # # 计算每个类别的精确率
- # precision = torch.diag(h) / torch.maximum(h.sum(0), torch.tensor(1))
- # # 计算F1分数
- # F1_Score = ((2 * precision(h) * recall(h)) / (precision(h) + recall(h))).sum()
-
- # 计算每个类别的召回率
- recall = np.diag(hist) / np.maximum(hist.sum(1), 1)
- # 计算每个类别的精确率
- precision = np.diag(hist) / np.maximum(hist.sum(0), 1)
- # 计算F1分数
- F1_Score = (2 * precision * recall) / (precision + recall)
- return acc_global, acc, iu, recall, precision, F1_Score
-
- def reduce_from_all_processes(self):
- if not torch.distributed.is_available():
- return
- if not torch.distributed.is_initialized():
- return
- torch.distributed.barrier()
- torch.distributed.all_reduce(self.mat)
-
- def __str__(self):
- acc_global, acc, iu, recall, precision, F1_Score = self.compute()
- return (
- 'global correct: {:.1f}\n'
- 'average row correct: {}\n'
- 'IoU: {}\n'
- 'recall: {}\n'
- 'precision: {}\n'
- 'F1_Score: {}\n'
- 'mean IoU: {:.1f}\n').format(
- acc_global.item() * 100,
- ['{:.1f}'.format(i) for i in (acc * 100).tolist()],
- ['{:.1f}'.format(i) for i in (iu * 100).tolist()],
- ['{:.1f}'.format(i) for i in (recall * 100).tolist()],
- ['{:.1f}'.format(i) for i in (precision * 100).tolist()],
- ['{:.1f}'.format(i) for i in (F1_Score * 100).tolist()],
- iu.mean().item() * 100)
-
-
- class MetricLogger(object):
- def __init__(self, delimiter="\t"):
- self.meters = defaultdict(SmoothedValue)
- self.delimiter = delimiter
-
- def update(self, **kwargs):
- for k, v in kwargs.items():
- if isinstance(v, torch.Tensor):
- v = v.item()
- assert isinstance(v, (float, int))
- self.meters[k].update(v)
-
- def __getattr__(self, attr):
- if attr in self.meters:
- return self.meters[attr]
- if attr in self.__dict__:
- return self.__dict__[attr]
- raise AttributeError("'{}' object has no attribute '{}'".format(
- type(self).__name__, attr))
-
- def __str__(self):
- loss_str = []
- for name, meter in self.meters.items():
- loss_str.append(
- "{}: {}".format(name, str(meter))
- )
- return self.delimiter.join(loss_str)
-
- def synchronize_between_processes(self):
- for meter in self.meters.values():
- meter.synchronize_between_processes()
-
- def add_meter(self, name, meter):
- self.meters[name] = meter
-
- def log_every(self, iterable, print_freq, header=None):
- i = 0
- if not header:
- header = ''
- start_time = time.time()
- end = time.time()
- iter_time = SmoothedValue(fmt='{avg:.4f}')
- data_time = SmoothedValue(fmt='{avg:.4f}')
- space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
- if torch.cuda.is_available():
- log_msg = self.delimiter.join([
- header,
- '[{0' + space_fmt + '}/{1}]',
- 'eta: {eta}',
- '{meters}',
- 'time: {time}',
- 'data: {data}',
- 'max mem: {memory:.0f}'
- ])
- else:
- log_msg = self.delimiter.join([
- header,
- '[{0' + space_fmt + '}/{1}]',
- 'eta: {eta}',
- '{meters}',
- 'time: {time}',
- 'data: {data}'
- ])
- MB = 256.0 * 256.0
- for obj in iterable:
- data_time.update(time.time() - end)
- yield obj
- iter_time.update(time.time() - end)
- if i % print_freq == 0:
- eta_seconds = iter_time.global_avg * (len(iterable) - i)
- eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
- if torch.cuda.is_available():
- print(log_msg.format(
- i, len(iterable), eta=eta_string,
- meters=str(self),
- time=str(iter_time), data=str(data_time),
- memory=torch.cuda.max_memory_allocated() / MB))
- else:
- print(log_msg.format(
- i, len(iterable), eta=eta_string,
- meters=str(self),
- time=str(iter_time), data=str(data_time)))
- i += 1
- end = time.time()
- total_time = time.time() - start_time
- total_time_str = str(datetime.timedelta(seconds=int(total_time)))
- print('{} Total time: {}'.format(header, total_time_str))
-
-
- def mkdir(path):
- try:
- os.makedirs(path)
- except OSError as e:
- if e.errno != errno.EEXIST:
- raise
-
-
- def setup_for_distributed(is_master):
- """
- This function disables printing when not in master process
- """
- import builtins as __builtin__
- builtin_print = __builtin__.print
-
- def print(*args, **kwargs):
- force = kwargs.pop('force', False)
- if is_master or force:
- builtin_print(*args, **kwargs)
-
- __builtin__.print = print
-
-
- def is_dist_avail_and_initialized():
- if not dist.is_available():
- return False
- if not dist.is_initialized():
- return False
- return True
-
-
- def get_world_size():
- if not is_dist_avail_and_initialized():
- return 1
- return dist.get_world_size()
-
-
- def get_rank():
- if not is_dist_avail_and_initialized():
- return 0
- return dist.get_rank()
-
-
- def is_main_process():
- return get_rank() == 0
-
-
- def save_on_master(*args, **kwargs):
- if is_main_process():
- torch.save(*args, **kwargs)
-
-
- def init_distributed_mode(args):
- if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
- args.rank = int(os.environ["RANK"])
- args.world_size = int(os.environ['WORLD_SIZE'])
- args.gpu = int(os.environ['LOCAL_RANK'])
- elif 'SLURM_PROCID' in os.environ:
- args.rank = int(os.environ['SLURM_PROCID'])
- args.gpu = args.rank % torch.cuda.device_count()
- elif hasattr(args, "rank"):
- pass
- else:
- print('Not using distributed mode')
- args.distributed = False
- return
-
- args.distributed = True
-
- torch.cuda.set_device(args.gpu)
- args.dist_backend = 'nccl'
- print('| distributed init (rank {}): {}'.format(
- args.rank, args.dist_url), flush=True)
- torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
- world_size=args.world_size, rank=args.rank)
- setup_for_distributed(args.rank == 0)
|