xdecoder/Instruct-X-Decoder

1# Copyright (c) Facebook, Inc. and its affiliates.2# Modified by Bowen Cheng from https://github.com/facebookresearch/detr/blob/master/util/misc.py3# Modified by Xueyan Zou4"""5Misc functions, including distributed helpers.6 7Mostly copy-paste from torchvision references.8"""9from typing import List, Optional10 11import torch12import torch.distributed as dist13import torchvision14from torch import Tensor15 16def _max_by_axis(the_list):17    # type: (List[List[int]]) -> List[int]18    maxes = the_list[0]19    for sublist in the_list[1:]:20        for index, item in enumerate(sublist):21            maxes[index] = max(maxes[index], item)22    return maxes23 24class NestedTensor(object):25    def __init__(self, tensors, mask: Optional[Tensor]):26        self.tensors = tensors27        self.mask = mask28 29    def to(self, device):30        # type: (Device) -> NestedTensor # noqa31        cast_tensor = self.tensors.to(device)32        mask = self.mask33        if mask is not None:34            assert mask is not None35            cast_mask = mask.to(device)36        else:37            cast_mask = None38        return NestedTensor(cast_tensor, cast_mask)39 40    def decompose(self):41        return self.tensors, self.mask42 43    def __repr__(self):44        return str(self.tensors)45 46def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):47    # TODO make this more general48    if tensor_list[0].ndim == 3:49        if torchvision._is_tracing():50            # nested_tensor_from_tensor_list() does not export well to ONNX51            # call _onnx_nested_tensor_from_tensor_list() instead52            return _onnx_nested_tensor_from_tensor_list(tensor_list)53 54        # TODO make it support different-sized images55        max_size = _max_by_axis([list(img.shape) for img in tensor_list])56        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))57        batch_shape = [len(tensor_list)] + max_size58        b, c, h, w = batch_shape59        dtype = tensor_list[0].dtype60        device = tensor_list[0].device61        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)62        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)63        for img, pad_img, m in zip(tensor_list, tensor, mask):64            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)65            m[: img.shape[1], : img.shape[2]] = False66    elif tensor_list[0].ndim == 2:67        if torchvision._is_tracing():68            # nested_tensor_from_tensor_list() does not export well to ONNX69            # call _onnx_nested_tensor_from_tensor_list() instead70            return _onnx_nested_tensor_from_tensor_list(tensor_list)71 72        # TODO make it support different-sized images73        max_size = _max_by_axis([list(txt.shape) for txt in tensor_list])74        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))75        batch_shape = [len(tensor_list)] + max_size76        b, c, l = batch_shape77        dtype = tensor_list[0].dtype78        device = tensor_list[0].device79        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)80        mask = torch.ones((b, l), dtype=torch.bool, device=device)81        for txt, pad_txt, m in zip(tensor_list, tensor, mask):82            pad_txt[: txt.shape[0], : txt.shape[1]] = txt83            m[: txt.shape[1]] = False84    else:85        raise ValueError("not supported")86    return NestedTensor(tensor, mask)87 88def _collate_and_pad_divisibility(tensor_list: list, div=32):89    max_size = []90    for i in range(tensor_list[0].dim()):91        max_size_i = torch.max(92            torch.tensor([img.shape[i] for img in tensor_list]).to(torch.float32)93        ).to(torch.int64)94        max_size.append(max_size_i)95    max_size = tuple(max_size)96 97    c,h,w = max_size98    pad_h = (div - h % div) if h % div != 0 else 099    pad_w = (div - w % div) if w % div != 0 else 0100    max_size = (c,h+pad_h,w+pad_w)101    102    # work around for103    # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)104    # m[: img.shape[1], :img.shape[2]] = False105    # which is not yet supported in onnx106    padded_imgs = []107    padded_masks = []108    for img in tensor_list:109        padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]110        padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))111        padded_imgs.append(padded_img)112 113        m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)114        padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)115        padded_masks.append(padded_mask.to(torch.bool))116    117    return padded_imgs118 119# _onnx_nested_tensor_from_tensor_list() is an implementation of120# nested_tensor_from_tensor_list() that is supported by ONNX tracing.121@torch.jit.unused122def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor:123    max_size = []124    for i in range(tensor_list[0].dim()):125        max_size_i = torch.max(126            torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)127        ).to(torch.int64)128        max_size.append(max_size_i)129    max_size = tuple(max_size)130 131    # work around for132    # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)133    # m[: img.shape[1], :img.shape[2]] = False134    # which is not yet supported in onnx135    padded_imgs = []136    padded_masks = []137    for img in tensor_list:138        padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]139        padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))140        padded_imgs.append(padded_img)141 142        m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)143        padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)144        padded_masks.append(padded_mask.to(torch.bool))145 146    tensor = torch.stack(padded_imgs)147    mask = torch.stack(padded_masks)148 149    return NestedTensor(tensor, mask=mask)150 151 152def is_dist_avail_and_initialized():153    if not dist.is_available():154        return False155    if not dist.is_initialized():156        return False157    return True