|
|
import logging |
|
|
from contextlib import contextmanager |
|
|
from functools import wraps |
|
|
|
|
|
import torch |
|
|
from mmcv.cnn.bricks.wrappers import obsolete_torch_version |
|
|
from torch.nn import functional as F |
|
|
|
|
|
TORCH_VERSION = tuple(int(x) for x in torch.__version__.split('.')[:2]) |
|
|
|
|
|
|
|
|
def is_lower_torch_version(version=(1, 10)): |
|
|
"""Check if the pytorch version is lower than "version.""" |
|
|
return obsolete_torch_version(TORCH_VERSION, version) |
|
|
|
|
|
|
|
|
@contextmanager |
|
|
def _ignore_torch_cuda_oom(): |
|
|
"""A context which ignores CUDA OOM exception from pytorch.""" |
|
|
try: |
|
|
yield |
|
|
except RuntimeError as e: |
|
|
if 'CUDA out of memory. ' in str(e): |
|
|
pass |
|
|
else: |
|
|
raise |
|
|
|
|
|
|
|
|
def retry_if_cuda_oom(func): |
|
|
"""Makes a function retry itself after encountering pytorch's CUDA OOM |
|
|
error. It will first retry after calling `torch.cuda.empty_cache()`. |
|
|
|
|
|
If that still fails, it will then retry by trying to convert inputs |
|
|
to CPUs. In this case, it expects the function to dispatch to CPU |
|
|
implementation. The return values may become CPU tensors as well |
|
|
and it's user's responsibility to convert it back to CUDA tensor |
|
|
if needed. |
|
|
|
|
|
Args: |
|
|
func: a stateless callable that takes tensor-like objects as arguments |
|
|
|
|
|
Returns: |
|
|
a callable which retries `func` if OOM is encountered. |
|
|
|
|
|
Examples: |
|
|
:: |
|
|
output = retry_if_cuda_oom(some_torch_function)(input1, input2) |
|
|
# output may be on CPU even if inputs are on GPU |
|
|
|
|
|
Note: |
|
|
1. When converting inputs to CPU, it will only |
|
|
look at each argument and check if it has `.device` |
|
|
and `.to` for conversion. Nested structures of tensors |
|
|
are not supported. |
|
|
|
|
|
2. Since the function might be called more than once, it has to be |
|
|
stateless. |
|
|
""" |
|
|
|
|
|
def maybe_to_cpu(x): |
|
|
try: |
|
|
like_gpu_tensor = x.device.type == 'cuda' and hasattr(x, 'to') |
|
|
except AttributeError: |
|
|
like_gpu_tensor = False |
|
|
if like_gpu_tensor: |
|
|
return x.to(device='cpu') |
|
|
else: |
|
|
return x |
|
|
|
|
|
@wraps(func) |
|
|
def wrapped(*args, **kwargs): |
|
|
with _ignore_torch_cuda_oom(): |
|
|
return func(*args, **kwargs) |
|
|
|
|
|
|
|
|
torch.cuda.empty_cache() |
|
|
with _ignore_torch_cuda_oom(): |
|
|
return func(*args, **kwargs) |
|
|
|
|
|
|
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
logger.info( |
|
|
'Attempting to copy inputs of {} to CPU due to CUDA OOM'.format( |
|
|
str(func)[0:5])) |
|
|
new_args = (maybe_to_cpu(x) for x in args) |
|
|
new_kwargs = {k: maybe_to_cpu(v) for k, v in kwargs.items()} |
|
|
return func(*new_args, **new_kwargs) |
|
|
|
|
|
return wrapped |
|
|
|
|
|
|
|
|
def sem_seg_postprocess(result, img_size, output_height, output_width): |
|
|
"""Return semantic segmentation predictions in the original resolution. |
|
|
|
|
|
The input images are often resized when entering semantic segmentor. |
|
|
Moreover, in same cases, they also padded inside segmentor to be |
|
|
divisible by maximum network stride. As a result, we often need |
|
|
the predictions of the segmentor in a different resolution from |
|
|
its inputs. |
|
|
|
|
|
Args: |
|
|
result (Tensor): semantic segmentation prediction logits. |
|
|
A tensor of shape (C, H, W), where C is the number of classes, |
|
|
and H, W are the height and width of the prediction. |
|
|
img_size (tuple): image size that segmentor is taking as input. |
|
|
output_height, output_width: the desired output resolution. |
|
|
|
|
|
Returns: |
|
|
semantic segmentation prediction (Tensor): A tensor of the shape |
|
|
(C, output_height, output_width) that contains per-pixel |
|
|
soft predictions. |
|
|
""" |
|
|
result = result[:, :img_size[0], :img_size[1]].expand(1, -1, -1, -1) |
|
|
if is_lower_torch_version(): |
|
|
result = F.interpolate( |
|
|
result, |
|
|
size=(output_height, output_width), |
|
|
mode='bicubic', |
|
|
align_corners=False)[0] |
|
|
else: |
|
|
result = F.interpolate( |
|
|
result, |
|
|
size=(output_height, output_width), |
|
|
mode='bicubic', |
|
|
align_corners=False, |
|
|
antialias=True)[0] |
|
|
return result |
|
|
|
|
|
|
|
|
def get_prompt_templates(): |
|
|
prompt_templates = [ |
|
|
'{}.', |
|
|
'a photo of a {}.', |
|
|
'a bad photo of a {}.', |
|
|
'a photo of many {}.', |
|
|
'a sculpture of a {}.', |
|
|
'a photo of the hard to see {}.', |
|
|
'a low resolution photo of the {}.', |
|
|
'a rendering of a {}.', |
|
|
'graffiti of a {}.', |
|
|
'a bad photo of the {}.', |
|
|
'a cropped photo of the {}.', |
|
|
'a tattoo of a {}.', |
|
|
'the embroidered {}.', |
|
|
'a photo of a hard to see {}.', |
|
|
'a bright photo of a {}.', |
|
|
'a photo of a clean {}.', |
|
|
'a photo of a dirty {}.', |
|
|
'a dark photo of the {}.', |
|
|
'a drawing of a {}.', |
|
|
'a photo of my {}.', |
|
|
'the plastic {}.', |
|
|
'a photo of the cool {}.', |
|
|
'a close-up photo of a {}.', |
|
|
'a black and white photo of the {}.', |
|
|
'a painting of the {}.', |
|
|
'a painting of a {}.', |
|
|
'a pixelated photo of the {}.', |
|
|
'a sculpture of the {}.', |
|
|
'a bright photo of the {}.', |
|
|
'a cropped photo of a {}.', |
|
|
'a plastic {}.', |
|
|
'a photo of the dirty {}.', |
|
|
'a jpeg corrupted photo of a {}.', |
|
|
'a blurry photo of the {}.', |
|
|
'a photo of the {}.', |
|
|
'a good photo of the {}.', |
|
|
'a rendering of the {}.', |
|
|
'a {} in a video game.', |
|
|
'a photo of one {}.', |
|
|
'a doodle of a {}.', |
|
|
'a close-up photo of the {}.', |
|
|
'the origami {}.', |
|
|
'the {} in a video game.', |
|
|
'a sketch of a {}.', |
|
|
'a doodle of the {}.', |
|
|
'a origami {}.', |
|
|
'a low resolution photo of a {}.', |
|
|
'the toy {}.', |
|
|
'a rendition of the {}.', |
|
|
'a photo of the clean {}.', |
|
|
'a photo of a large {}.', |
|
|
'a rendition of a {}.', |
|
|
'a photo of a nice {}.', |
|
|
'a photo of a weird {}.', |
|
|
'a blurry photo of a {}.', |
|
|
'a cartoon {}.', |
|
|
'art of a {}.', |
|
|
'a sketch of the {}.', |
|
|
'a embroidered {}.', |
|
|
'a pixelated photo of a {}.', |
|
|
'itap of the {}.', |
|
|
'a jpeg corrupted photo of the {}.', |
|
|
'a good photo of a {}.', |
|
|
'a plushie {}.', |
|
|
'a photo of the nice {}.', |
|
|
'a photo of the small {}.', |
|
|
'a photo of the weird {}.', |
|
|
'the cartoon {}.', |
|
|
'art of the {}.', |
|
|
'a drawing of the {}.', |
|
|
'a photo of the large {}.', |
|
|
'a black and white photo of a {}.', |
|
|
'the plushie {}.', |
|
|
'a dark photo of a {}.', |
|
|
'itap of a {}.', |
|
|
'graffiti of the {}.', |
|
|
'a toy {}.', |
|
|
'itap of my {}.', |
|
|
'a photo of a cool {}.', |
|
|
'a photo of a small {}.', |
|
|
'a tattoo of the {}.', |
|
|
] |
|
|
return prompt_templates |
|
|
|
