| from enum import Enum |
|
|
| import numpy as np |
| import torch |
| import torch.distributed as dist |
|
|
| IGNORE_INDEX = -100 |
| IMAGE_TOKEN_INDEX = -200 |
|
|
| DEFAULT_EOS_TOKEN = "</s>" |
| DEFAULT_BOS_TOKEN = "<s>" |
| DEFAULT_UNK_TOKEN = "<unk>" |
|
|
| DEFAULT_IMAGE_TOKEN = "<image>" |
| DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>" |
| DEFAULT_IM_START_TOKEN = "<im_start>" |
| DEFAULT_IM_END_TOKEN = "<im_end>" |
| DEFAULT_BBOX_TOKEN = "<bbox>" |
|
|
|
|
| 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 to_cuda(e): |
| return ( |
| e.cuda(non_blocking=True) |
| if e is not None and not isinstance(e, (str, int)) |
| else e |
| ) |
|
|
|
|
| def dict_to_cuda(input_dict): |
| for k, v in input_dict.items(): |
| if isinstance(v, torch.Tensor): |
| input_dict[k] = to_cuda(v) |
| elif isinstance(input_dict[k], list): |
| input_dict[k] = list_to_cuda(input_dict[k]) |
| elif isinstance(input_dict[k], tuple): |
| input_dict[k] = tuple_to_cuda(input_dict[k]) |
| elif isinstance(v, dict): |
| input_dict[k] = dict_to_cuda(v) |
| return input_dict |
|
|
|
|
| def list_to_cuda(input_list): |
| for i in range(len(input_list)): |
| if isinstance(input_list[i], torch.Tensor): |
| input_list[i] = to_cuda(input_list[i]) |
| elif isinstance(input_list[i], dict): |
| input_list[i] = dict_to_cuda(input_list[i]) |
| elif isinstance(input_list[i], list): |
| input_list[i] = list_to_cuda(input_list[i]) |
| elif isinstance(input_list[i], tuple): |
| input_list[i] = tuple_to_cuda(input_list[i]) |
| return input_list |
|
|
|
|
| def tuple_to_cuda(input_tuple): |
| return tuple( |
| ( |
| to_cuda(e) |
| if isinstance(e, torch.Tensor) |
| else ( |
| dict_to_cuda(e) |
| if isinstance(e, dict) |
| else ( |
| list_to_cuda(e) |
| if isinstance(e, list) |
| else tuple_to_cuda(e) if isinstance(e, tuple) else e |
| ) |
| ) |
| ) |
| for e in input_tuple |
| ) |
|
|
