|
|
|
|
|
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 |
