|
|
|
|
|
import warnings |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
|
|
|
from .registry import MODELS |
|
|
|
|
|
|
|
|
TORCH_VERSION = torch.__version__ |
|
|
|
|
|
def resize(input, |
|
|
size=None, |
|
|
scale_factor=None, |
|
|
mode='nearest', |
|
|
align_corners=None, |
|
|
warning=True): |
|
|
if warning: |
|
|
if size is not None and align_corners: |
|
|
input_h, input_w = tuple(int(x) for x in input.shape[2:]) |
|
|
output_h, output_w = tuple(int(x) for x in size) |
|
|
if output_h > input_h or output_w > output_h: |
|
|
if ((output_h > 1 and output_w > 1 and input_h > 1 |
|
|
and input_w > 1) and (output_h - 1) % (input_h - 1) |
|
|
and (output_w - 1) % (input_w - 1)): |
|
|
warnings.warn( |
|
|
f'When align_corners={align_corners}, ' |
|
|
'the output would more aligned if ' |
|
|
f'input size {(input_h, input_w)} is `x+1` and ' |
|
|
f'out size {(output_h, output_w)} is `nx+1`') |
|
|
return F.interpolate(input, size, scale_factor, mode, align_corners) |
|
|
|
|
|
|
|
|
@MODELS.register_module('Conv', force=True) |
|
|
class Conv2d(nn.Conv2d): |
|
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
|
if x.numel() == 0 and obsolete_torch_version(TORCH_VERSION, (1, 4)): |
|
|
out_shape = [x.shape[0], self.out_channels] |
|
|
for i, k, p, s, d in zip(x.shape[-2:], self.kernel_size, |
|
|
self.padding, self.stride, self.dilation): |
|
|
o = (i + 2 * p - (d * (k - 1) + 1)) // s + 1 |
|
|
out_shape.append(o) |
|
|
empty = NewEmptyTensorOp.apply(x, out_shape) |
|
|
if self.training: |
|
|
|
|
|
dummy = sum(x.view(-1)[0] for x in self.parameters()) * 0.0 |
|
|
return empty + dummy |
|
|
else: |
|
|
return empty |
|
|
|
|
|
return super().forward(x) |
|
|
|
|
|
|
|
|
class NewEmptyTensorOp(torch.autograd.Function): |
|
|
|
|
|
@staticmethod |
|
|
def forward(ctx, x: torch.Tensor, new_shape: tuple) -> torch.Tensor: |
|
|
ctx.shape = x.shape |
|
|
return x.new_empty(new_shape) |
|
|
|
|
|
@staticmethod |
|
|
def backward(ctx, grad: torch.Tensor) -> tuple: |
|
|
shape = ctx.shape |
|
|
return NewEmptyTensorOp.apply(grad, shape), None |
|
|
|
|
|
|
|
|
def obsolete_torch_version(torch_version, version_threshold) -> bool: |
|
|
return torch_version == 'parrots' or torch_version <= version_threshold |
|
|
|
|
|
|
|
|
@MODELS.register_module() |
|
|
class DropPath(nn.Module): |
|
|
"""Drop paths (Stochastic Depth) per sample (when applied in main path of |
|
|
residual blocks). |
|
|
|
|
|
We follow the implementation |
|
|
https://github.com/rwightman/pytorch-image-models/blob/a2727c1bf78ba0d7b5727f5f95e37fb7f8866b1f/timm/models/layers/drop.py # noqa: E501 |
|
|
|
|
|
Args: |
|
|
drop_prob (float): Probability of the path to be zeroed. Default: 0.1 |
|
|
""" |
|
|
|
|
|
def __init__(self, drop_prob: float = 0.1): |
|
|
super().__init__() |
|
|
self.drop_prob = drop_prob |
|
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
|
return drop_path(x, self.drop_prob, self.training) |
|
|
|
|
|
|
|
|
def drop_path(x: torch.Tensor, |
|
|
drop_prob: float = 0., |
|
|
training: bool = False) -> torch.Tensor: |
|
|
"""Drop paths (Stochastic Depth) per sample (when applied in main path of |
|
|
residual blocks). |
|
|
|
|
|
We follow the implementation |
|
|
https://github.com/rwightman/pytorch-image-models/blob/a2727c1bf78ba0d7b5727f5f95e37fb7f8866b1f/timm/models/layers/drop.py # noqa: E501 |
|
|
""" |
|
|
if drop_prob == 0. or not training: |
|
|
return x |
|
|
keep_prob = 1 - drop_prob |
|
|
|
|
|
shape = (x.shape[0], ) + (1, ) * (x.ndim - 1) |
|
|
random_tensor = keep_prob + torch.rand( |
|
|
shape, dtype=x.dtype, device=x.device) |
|
|
output = x.div(keep_prob) * random_tensor.floor() |
|
|
return output |
