Pikaliov
71862bbeca
- Add dcnv4_int8_cuda.cu/.h: CUDA kernel (int8 values, fp16 offsets, fp32 interp, requantize with value_scale/output_scale) - Add dcnv4_int8_forward(): inference-only functional wrapper (@no_grad) - Add DCNv4Strip.forward_int8(): module-level INT8 forward (without_pointwise=True) - Add scripts/test_dcnv4_int8.py: correctness gate (<=1 LSB, >=99% exact) and informational fp16 vs int8 benchmark - Update README: INT8 API section, updated structure tree, SOFIA/MERIDIAN context Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
325 lines
13 KiB
Python
325 lines
13 KiB
Python
# --------------------------------------------------------
|
||
# Deformable Convolution v4
|
||
# Copyright (c) 2023 OpenGVLab
|
||
# Licensed under The MIT License [see LICENSE for details]
|
||
# --------------------------------------------------------
|
||
|
||
from __future__ import absolute_import
|
||
from __future__ import print_function
|
||
from __future__ import division
|
||
|
||
import math
|
||
import torch
|
||
from torch import nn
|
||
import torch.nn.functional as F
|
||
from torch.nn.init import xavier_uniform_, constant_
|
||
from ..functions import DCNv4Function, dcnv4_int8_forward
|
||
|
||
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 DCNv4(nn.Module):
|
||
def __init__(
|
||
self,
|
||
channels=64,
|
||
kernel_size=3,
|
||
stride=1,
|
||
pad=1,
|
||
dilation=1,
|
||
group=4,
|
||
offset_scale=1.0,
|
||
dw_kernel_size=None,
|
||
center_feature_scale=False,
|
||
remove_center=False,
|
||
output_bias=True,
|
||
without_pointwise=False,
|
||
**kwargs):
|
||
"""
|
||
DCNv4 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
|
||
|
||
# you'd better set _d_per_group to a power of 2 which is more efficient in our CUDA implementation
|
||
assert _d_per_group % 16 == 0
|
||
|
||
self.offset_scale = offset_scale
|
||
self.channels = channels
|
||
self.kernel_size = 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.dw_kernel_size = dw_kernel_size
|
||
self.center_feature_scale = center_feature_scale
|
||
self.remove_center = int(remove_center)
|
||
self.without_pointwise = without_pointwise
|
||
|
||
self.K = group * (kernel_size * kernel_size - self.remove_center)
|
||
if dw_kernel_size is not None:
|
||
self.offset_mask_dw = nn.Conv2d(channels, channels, dw_kernel_size, stride=1, padding=(dw_kernel_size - 1) // 2, groups=channels)
|
||
self.offset_mask = nn.Linear(channels, int(math.ceil((self.K * 3)/8)*8))
|
||
if not without_pointwise:
|
||
self.value_proj = nn.Linear(channels, channels)
|
||
self.output_proj = nn.Linear(channels, channels, bias=output_bias)
|
||
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_mask.weight.data, 0.)
|
||
constant_(self.offset_mask.bias.data, 0.)
|
||
if not self.without_pointwise:
|
||
xavier_uniform_(self.value_proj.weight.data)
|
||
constant_(self.value_proj.bias.data, 0.)
|
||
xavier_uniform_(self.output_proj.weight.data)
|
||
if self.output_proj.bias is not None:
|
||
constant_(self.output_proj.bias.data, 0.)
|
||
|
||
def forward(self, input, shape=None):
|
||
"""
|
||
:param query (N, H, W, C)
|
||
:return output (N, H, W, C)
|
||
"""
|
||
N, L, C = input.shape
|
||
if shape is not None:
|
||
H, W = shape
|
||
else:
|
||
H, W = int(L**0.5), int(L**0.5)
|
||
|
||
|
||
x = input
|
||
if not self.without_pointwise:
|
||
x = self.value_proj(x)
|
||
x = x.reshape(N, H, W, -1)
|
||
if self.dw_kernel_size is not None:
|
||
offset_mask_input = self.offset_mask_dw(input.view(N, H, W, C).permute(0, 3, 1, 2))
|
||
offset_mask_input = offset_mask_input.permute(0, 2, 3, 1).view(N, L, C)
|
||
else:
|
||
offset_mask_input = input
|
||
offset_mask = self.offset_mask(offset_mask_input).reshape(N, H, W, -1)
|
||
|
||
x_proj = x
|
||
|
||
x = DCNv4Function.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(
|
||
x, 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 = x.view(N, L, -1)
|
||
|
||
if not self.without_pointwise:
|
||
x = self.output_proj(x)
|
||
return x
|
||
|
||
|
||
# Kernel-point counts (kernel_h * kernel_w) that have compiled template
|
||
# instantiations in dcnv4_im2col_cuda.cuh / dcnv4_col2im_cuda.cuh
|
||
# (``switch (K) { case 9 / 25 / 49 }``). Any other K requires adding a case
|
||
# to those switches and rebuilding the extension.
|
||
_COMPILED_K = (9, 25, 49)
|
||
|
||
|
||
class DCNv4Strip(nn.Module):
|
||
"""Deformable STRIP convolution: a (1, k) or (k, 1) DCNv4 sampling line.
|
||
|
||
SOFIA "strip-DCN" (O2 in RECOMMENDATIONS Дополнение 3): the deformable
|
||
neighbourhood is a line of ``k`` points instead of a k×k square, so the
|
||
offset/mask predictor shrinks by ~k× (e.g. K: 49 → 9 per group) while the
|
||
receptive field along the strip is preserved and offsets let the line bend
|
||
along image structures (roads, field boundaries).
|
||
|
||
The CUDA kernels are generic over (kernel_h, kernel_w) at runtime but
|
||
template-dispatch on K = kernel_h * kernel_w with compiled cases
|
||
{9, 25, 49} — therefore ``k`` defaults to 9 (a (1, 9) strip reuses the
|
||
existing K=9 instantiation; no rebuild needed). For other ``k`` extend the
|
||
switch in the .cuh files first.
|
||
|
||
Args:
|
||
channels: Input/output channels (sequence layout [N, L, C]).
|
||
k: Number of sampling points along the strip. Must be in {9, 25, 49}
|
||
unless ``allow_uncompiled_k=True`` (then you must have rebuilt the
|
||
extension with the extra case).
|
||
orientation: 'h' → kernel (1, k); 'v' → kernel (k, 1).
|
||
group: Offset groups; ``channels // group`` must be divisible by 16.
|
||
offset_scale: DCNv4 offset scale.
|
||
without_pointwise: Skip value/output projections (default True — in
|
||
SOFIA the surrounding MBConv 1×1s already mix channels).
|
||
output_bias: Bias for output projection (used only if pointwise on).
|
||
allow_uncompiled_k: Permit k outside the compiled set (see above).
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
channels=64,
|
||
k=9,
|
||
orientation='h',
|
||
group=4,
|
||
offset_scale=1.0,
|
||
without_pointwise=True,
|
||
output_bias=True,
|
||
allow_uncompiled_k=False,
|
||
**kwargs):
|
||
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
|
||
assert _d_per_group % 16 == 0, (
|
||
f'channels // group must be divisible by 16, got {_d_per_group}')
|
||
if orientation not in ('h', 'v'):
|
||
raise ValueError(f"orientation must be 'h' or 'v', got {orientation!r}")
|
||
if k not in _COMPILED_K and not allow_uncompiled_k:
|
||
raise ValueError(
|
||
f'k={k} has no compiled CUDA instantiation (K must be in '
|
||
f'{_COMPILED_K}); add a `case {k}:` to dcnv4_im2col_cuda.cuh / '
|
||
f'dcnv4_col2im_cuda.cuh and rebuild, then pass '
|
||
f'allow_uncompiled_k=True')
|
||
|
||
self.channels = channels
|
||
self.k = k
|
||
self.orientation = orientation
|
||
if orientation == 'h':
|
||
self.kernel_h, self.kernel_w = 1, k
|
||
self.pad_h, self.pad_w = 0, (k - 1) // 2
|
||
else:
|
||
self.kernel_h, self.kernel_w = k, 1
|
||
self.pad_h, self.pad_w = (k - 1) // 2, 0
|
||
self.stride = 1
|
||
self.dilation = 1
|
||
self.group = group
|
||
self.group_channels = channels // group
|
||
self.offset_scale = offset_scale
|
||
self.without_pointwise = without_pointwise
|
||
|
||
# Total points across groups; offsets (2K) + masks (K), padded to /8.
|
||
self.K = group * k
|
||
self.offset_mask = nn.Linear(channels, int(math.ceil((self.K * 3) / 8) * 8))
|
||
if not without_pointwise:
|
||
self.value_proj = nn.Linear(channels, channels)
|
||
self.output_proj = nn.Linear(channels, channels, bias=output_bias)
|
||
self._reset_parameters()
|
||
|
||
def _reset_parameters(self):
|
||
# Zero-init offsets/masks: the strip starts as a uniform static line
|
||
# (stable start; masks=0 → zero branch output, residual/other branches
|
||
# carry the signal until offsets learn).
|
||
constant_(self.offset_mask.weight.data, 0.)
|
||
constant_(self.offset_mask.bias.data, 0.)
|
||
if not self.without_pointwise:
|
||
xavier_uniform_(self.value_proj.weight.data)
|
||
constant_(self.value_proj.bias.data, 0.)
|
||
xavier_uniform_(self.output_proj.weight.data)
|
||
if self.output_proj.bias is not None:
|
||
constant_(self.output_proj.bias.data, 0.)
|
||
|
||
def forward(self, input, shape=None):
|
||
"""input: [N, L, C] (sequence layout, as DCNv4); returns [N, L, C]."""
|
||
N, L, C = input.shape
|
||
if shape is not None:
|
||
H, W = shape
|
||
else:
|
||
H, W = int(L ** 0.5), int(L ** 0.5)
|
||
|
||
x = input
|
||
if not self.without_pointwise:
|
||
x = self.value_proj(x)
|
||
x = x.reshape(N, H, W, -1)
|
||
offset_mask = self.offset_mask(input).reshape(N, H, W, -1)
|
||
|
||
x = DCNv4Function.apply(
|
||
x, offset_mask,
|
||
self.kernel_h, self.kernel_w,
|
||
self.stride, self.stride,
|
||
self.pad_h, self.pad_w,
|
||
self.dilation, self.dilation,
|
||
self.group, self.group_channels,
|
||
self.offset_scale,
|
||
256,
|
||
0, # remove_center: keep the centre point of the strip
|
||
)
|
||
|
||
x = x.view(N, L, -1)
|
||
if not self.without_pointwise:
|
||
x = self.output_proj(x)
|
||
return x
|
||
|
||
@torch.no_grad()
|
||
def forward_int8(self, input_int8, shape, value_scale, output_scale):
|
||
"""INT8 inference forward (Level 1: int8 storage + fp32 math).
|
||
|
||
The deformable sampling consumes/produces int8; the (small) offset/mask
|
||
Linear runs in the module's own float dtype on the dequantized input.
|
||
Requires ``without_pointwise=True`` (the int8 path quantizes the
|
||
sampling op only; pointwise projections belong to the surrounding
|
||
block's quantized 1x1 convs).
|
||
|
||
Args:
|
||
input_int8: int8 [N, L, C]; real value = input_int8 * value_scale.
|
||
shape: (H, W) of the token grid.
|
||
value_scale: per-tensor input scale.
|
||
output_scale: per-tensor output scale.
|
||
|
||
Returns:
|
||
int8 [N, L, C]; real output = result * output_scale.
|
||
"""
|
||
assert self.without_pointwise, \
|
||
"forward_int8 supports without_pointwise=True only"
|
||
assert input_int8.dtype == torch.int8, "input must be int8"
|
||
N, L, C = input_int8.shape
|
||
H, W = shape
|
||
|
||
w_dtype = self.offset_mask.weight.dtype
|
||
x_deq = input_int8.to(w_dtype) * float(value_scale)
|
||
offset_mask = self.offset_mask(x_deq).to(torch.float16)
|
||
offset_mask = offset_mask.reshape(N, H, W, -1).contiguous()
|
||
|
||
out = dcnv4_int8_forward(
|
||
input_int8.reshape(N, H, W, C).contiguous(), offset_mask,
|
||
self.kernel_h, self.kernel_w, self.pad_h, self.pad_w,
|
||
self.group, self.group_channels,
|
||
self.offset_scale, value_scale, output_scale,
|
||
)
|
||
return out.reshape(N, L, C)
|
||
|