|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from __future__ import absolute_import, division, print_function |
|
|
|
|
|
import warnings |
|
|
|
|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
from torch import nn |
|
|
from torch.nn.init import constant_, xavier_uniform_ |
|
|
|
|
|
from .dcnv3_func import DCNv3Function, dcnv3_core_pytorch, has_cuda_kernel |
|
|
|
|
|
|
|
|
class to_channels_first(nn.Module): |
|
|
|
|
|
def __init__(self): |
|
|
super().__init__() |
|
|
|
|
|
def forward(self, x): |
|
|
return x.permute(0, 3, 1, 2) |
|
|
|
|
|
|
|
|
class to_channels_last(nn.Module): |
|
|
|
|
|
def __init__(self): |
|
|
super().__init__() |
|
|
|
|
|
def forward(self, x): |
|
|
return x.permute(0, 2, 3, 1) |
|
|
|
|
|
|
|
|
def build_norm_layer(dim, |
|
|
norm_layer, |
|
|
in_format='channels_last', |
|
|
out_format='channels_last', |
|
|
eps=1e-6): |
|
|
layers = [] |
|
|
if norm_layer == 'BN': |
|
|
if in_format == 'channels_last': |
|
|
layers.append(to_channels_first()) |
|
|
layers.append(nn.BatchNorm2d(dim)) |
|
|
if out_format == 'channels_last': |
|
|
layers.append(to_channels_last()) |
|
|
elif norm_layer == 'LN': |
|
|
if in_format == 'channels_first': |
|
|
layers.append(to_channels_last()) |
|
|
layers.append(nn.LayerNorm(dim, eps=eps)) |
|
|
if out_format == 'channels_first': |
|
|
layers.append(to_channels_first()) |
|
|
else: |
|
|
raise NotImplementedError( |
|
|
f'build_norm_layer does not support {norm_layer}') |
|
|
return nn.Sequential(*layers) |
|
|
|
|
|
|
|
|
def build_act_layer(act_layer): |
|
|
if act_layer == 'ReLU': |
|
|
return nn.ReLU(inplace=True) |
|
|
elif act_layer == 'SiLU': |
|
|
return nn.SiLU(inplace=True) |
|
|
elif act_layer == 'GELU': |
|
|
return nn.GELU() |
|
|
|
|
|
raise NotImplementedError(f'build_act_layer does not support {act_layer}') |
|
|
|
|
|
|
|
|
def _is_power_of_2(n): |
|
|
if (not isinstance(n, int)) or (n < 0): |
|
|
raise ValueError( |
|
|
'invalid input for _is_power_of_2: {} (type: {})'.format(n, type(n))) |
|
|
|
|
|
return (n & (n - 1) == 0) and n != 0 |
|
|
|
|
|
|
|
|
class CenterFeatureScaleModule(nn.Module): |
|
|
def forward(self, |
|
|
query, |
|
|
center_feature_scale_proj_weight, |
|
|
center_feature_scale_proj_bias): |
|
|
center_feature_scale = F.linear(query, |
|
|
weight=center_feature_scale_proj_weight, |
|
|
bias=center_feature_scale_proj_bias).sigmoid() |
|
|
return center_feature_scale |
|
|
|
|
|
|
|
|
class DCNv3_pytorch(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
channels=64, |
|
|
kernel_size=3, |
|
|
dw_kernel_size=None, |
|
|
stride=1, |
|
|
pad=1, |
|
|
dilation=1, |
|
|
group=4, |
|
|
offset_scale=1.0, |
|
|
act_layer='GELU', |
|
|
norm_layer='LN', |
|
|
center_feature_scale=False, |
|
|
remove_center=False, |
|
|
): |
|
|
""" |
|
|
DCNv3 Module |
|
|
:param channels |
|
|
:param kernel_size |
|
|
:param stride |
|
|
:param pad |
|
|
:param dilation |
|
|
:param group |
|
|
:param offset_scale |
|
|
:param act_layer |
|
|
:param norm_layer |
|
|
""" |
|
|
super().__init__() |
|
|
if channels % group != 0: |
|
|
raise ValueError( |
|
|
f'channels must be divisible by group, but got {channels} and {group}') |
|
|
_d_per_group = channels // group |
|
|
dw_kernel_size = dw_kernel_size if dw_kernel_size is not None else kernel_size |
|
|
|
|
|
if not _is_power_of_2(_d_per_group): |
|
|
warnings.warn( |
|
|
"You'd better set channels in DCNv3 to make the dimension of each attention head a power of 2 " |
|
|
'which is more efficient in our CUDA implementation.') |
|
|
|
|
|
self.offset_scale = offset_scale |
|
|
self.channels = channels |
|
|
self.kernel_size = kernel_size |
|
|
self.dw_kernel_size = dw_kernel_size |
|
|
self.stride = stride |
|
|
self.dilation = dilation |
|
|
self.pad = pad |
|
|
self.group = group |
|
|
self.group_channels = channels // group |
|
|
self.offset_scale = offset_scale |
|
|
self.center_feature_scale = center_feature_scale |
|
|
self.remove_center = int(remove_center) |
|
|
|
|
|
self.dw_conv = nn.Sequential( |
|
|
nn.Conv2d( |
|
|
channels, |
|
|
channels, |
|
|
kernel_size=dw_kernel_size, |
|
|
stride=1, |
|
|
padding=(dw_kernel_size - 1) // 2, |
|
|
groups=channels), |
|
|
build_norm_layer( |
|
|
channels, |
|
|
norm_layer, |
|
|
'channels_first', |
|
|
'channels_last'), |
|
|
build_act_layer(act_layer)) |
|
|
self.offset = nn.Linear( |
|
|
channels, |
|
|
group * (kernel_size * kernel_size - remove_center) * 2) |
|
|
self.mask = nn.Linear( |
|
|
channels, |
|
|
group * (kernel_size * kernel_size - remove_center)) |
|
|
self.input_proj = nn.Linear(channels, channels) |
|
|
self.output_proj = nn.Linear(channels, channels) |
|
|
self._reset_parameters() |
|
|
|
|
|
if center_feature_scale: |
|
|
self.center_feature_scale_proj_weight = nn.Parameter( |
|
|
torch.zeros((group, channels), dtype=torch.float)) |
|
|
self.center_feature_scale_proj_bias = nn.Parameter( |
|
|
torch.tensor(0.0, dtype=torch.float).view((1,)).repeat(group, )) |
|
|
self.center_feature_scale_module = CenterFeatureScaleModule() |
|
|
|
|
|
def _reset_parameters(self): |
|
|
constant_(self.offset.weight.data, 0.) |
|
|
constant_(self.offset.bias.data, 0.) |
|
|
constant_(self.mask.weight.data, 0.) |
|
|
constant_(self.mask.bias.data, 0.) |
|
|
xavier_uniform_(self.input_proj.weight.data) |
|
|
constant_(self.input_proj.bias.data, 0.) |
|
|
xavier_uniform_(self.output_proj.weight.data) |
|
|
constant_(self.output_proj.bias.data, 0.) |
|
|
|
|
|
def forward(self, input): |
|
|
""" |
|
|
:param query (N, H, W, C) |
|
|
:return output (N, H, W, C) |
|
|
""" |
|
|
N, H, W, _ = input.shape |
|
|
|
|
|
x = self.input_proj(input) |
|
|
x_proj = x |
|
|
|
|
|
x1 = input.permute(0, 3, 1, 2) |
|
|
x1 = self.dw_conv(x1) |
|
|
offset = self.offset(x1) |
|
|
mask = self.mask(x1).reshape(N, H, W, self.group, -1) |
|
|
mask = F.softmax(mask, -1).reshape(N, H, W, -1) |
|
|
|
|
|
x = dcnv3_core_pytorch( |
|
|
x, offset, mask, |
|
|
self.kernel_size, self.kernel_size, |
|
|
self.stride, self.stride, |
|
|
self.pad, self.pad, |
|
|
self.dilation, self.dilation, |
|
|
self.group, self.group_channels, |
|
|
self.offset_scale, self.remove_center) |
|
|
if self.center_feature_scale: |
|
|
center_feature_scale = self.center_feature_scale_module( |
|
|
x1, self.center_feature_scale_proj_weight, self.center_feature_scale_proj_bias) |
|
|
|
|
|
center_feature_scale = center_feature_scale[..., None].repeat( |
|
|
1, 1, 1, 1, self.channels // self.group).flatten(-2) |
|
|
x = x * (1 - center_feature_scale) + x_proj * center_feature_scale |
|
|
x = self.output_proj(x) |
|
|
|
|
|
return x |
|
|
|
|
|
|
|
|
class DCNv3(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
channels=64, |
|
|
kernel_size=3, |
|
|
dw_kernel_size=None, |
|
|
stride=1, |
|
|
pad=1, |
|
|
dilation=1, |
|
|
group=4, |
|
|
offset_scale=1.0, |
|
|
act_layer='GELU', |
|
|
norm_layer='LN', |
|
|
center_feature_scale=False, |
|
|
remove_center=False, |
|
|
): |
|
|
""" |
|
|
DCNv3 Module |
|
|
:param channels |
|
|
:param kernel_size |
|
|
:param stride |
|
|
:param pad |
|
|
:param dilation |
|
|
:param group |
|
|
:param offset_scale |
|
|
:param act_layer |
|
|
:param norm_layer |
|
|
""" |
|
|
super().__init__() |
|
|
if channels % group != 0: |
|
|
raise ValueError( |
|
|
f'channels must be divisible by group, but got {channels} and {group}') |
|
|
_d_per_group = channels // group |
|
|
dw_kernel_size = dw_kernel_size if dw_kernel_size is not None else kernel_size |
|
|
|
|
|
if not _is_power_of_2(_d_per_group): |
|
|
warnings.warn( |
|
|
"You'd better set channels in DCNv3 to make the dimension of each attention head a power of 2 " |
|
|
'which is more efficient in our CUDA implementation.') |
|
|
|
|
|
self.offset_scale = offset_scale |
|
|
self.channels = channels |
|
|
self.kernel_size = kernel_size |
|
|
self.dw_kernel_size = dw_kernel_size |
|
|
self.stride = stride |
|
|
self.dilation = dilation |
|
|
self.pad = pad |
|
|
self.group = group |
|
|
self.group_channels = channels // group |
|
|
self.offset_scale = offset_scale |
|
|
self.center_feature_scale = center_feature_scale |
|
|
self.remove_center = int(remove_center) |
|
|
|
|
|
if self.remove_center and self.kernel_size % 2 == 0: |
|
|
raise ValueError('remove_center is only compatible with odd kernel size.') |
|
|
|
|
|
self.dw_conv = nn.Sequential( |
|
|
nn.Conv2d( |
|
|
channels, |
|
|
channels, |
|
|
kernel_size=dw_kernel_size, |
|
|
stride=1, |
|
|
padding=(dw_kernel_size - 1) // 2, |
|
|
groups=channels), |
|
|
build_norm_layer( |
|
|
channels, |
|
|
norm_layer, |
|
|
'channels_first', |
|
|
'channels_last'), |
|
|
build_act_layer(act_layer)) |
|
|
self.offset = nn.Linear( |
|
|
channels, |
|
|
group * (kernel_size * kernel_size - remove_center) * 2) |
|
|
self.mask = nn.Linear( |
|
|
channels, |
|
|
group * (kernel_size * kernel_size - remove_center)) |
|
|
self.input_proj = nn.Linear(channels, channels) |
|
|
self.output_proj = nn.Linear(channels, channels) |
|
|
self._reset_parameters() |
|
|
|
|
|
if center_feature_scale: |
|
|
self.center_feature_scale_proj_weight = nn.Parameter( |
|
|
torch.zeros((group, channels), dtype=torch.float)) |
|
|
self.center_feature_scale_proj_bias = nn.Parameter( |
|
|
torch.tensor(0.0, dtype=torch.float).view((1,)).repeat(group, )) |
|
|
self.center_feature_scale_module = CenterFeatureScaleModule() |
|
|
|
|
|
def _reset_parameters(self): |
|
|
constant_(self.offset.weight.data, 0.) |
|
|
constant_(self.offset.bias.data, 0.) |
|
|
constant_(self.mask.weight.data, 0.) |
|
|
constant_(self.mask.bias.data, 0.) |
|
|
xavier_uniform_(self.input_proj.weight.data) |
|
|
constant_(self.input_proj.bias.data, 0.) |
|
|
xavier_uniform_(self.output_proj.weight.data) |
|
|
constant_(self.output_proj.bias.data, 0.) |
|
|
|
|
|
def forward(self, input): |
|
|
""" |
|
|
:param query (N, H, W, C) |
|
|
:return output (N, H, W, C) |
|
|
""" |
|
|
N, H, W, _ = input.shape |
|
|
|
|
|
x = self.input_proj(input) |
|
|
x_proj = x |
|
|
dtype = x.dtype |
|
|
|
|
|
x1 = input.permute(0, 3, 1, 2) |
|
|
x1 = self.dw_conv(x1) |
|
|
offset = self.offset(x1) |
|
|
mask = self.mask(x1).reshape(N, H, W, self.group, -1) |
|
|
mask = F.softmax(mask, -1) |
|
|
mask = mask.reshape(N, H, W, -1).type(dtype) |
|
|
|
|
|
x = DCNv3Function.apply( |
|
|
x, offset, mask, |
|
|
self.kernel_size, self.kernel_size, |
|
|
self.stride, self.stride, |
|
|
self.pad, self.pad, |
|
|
self.dilation, self.dilation, |
|
|
self.group, self.group_channels, |
|
|
self.offset_scale, |
|
|
256, |
|
|
self.remove_center) |
|
|
|
|
|
if self.center_feature_scale: |
|
|
center_feature_scale = self.center_feature_scale_module( |
|
|
x1, self.center_feature_scale_proj_weight, self.center_feature_scale_proj_bias) |
|
|
|
|
|
center_feature_scale = center_feature_scale[..., None].repeat( |
|
|
1, 1, 1, 1, self.channels // self.group).flatten(-2) |
|
|
x = x * (1 - center_feature_scale) + x_proj * center_feature_scale |
|
|
x = self.output_proj(x) |
|
|
|
|
|
return x |
|
|
|
