|
|
"""
|
|
|
https://github.com/PaddlePaddle/PaddleDetection/blob/release/2.6/ppdet/modeling/backbones/cspresnet.py
|
|
|
|
|
|
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
|
|
|
"""
|
|
|
|
|
|
from collections import OrderedDict
|
|
|
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
import torch.nn.functional as F
|
|
|
|
|
|
from ...core import register
|
|
|
from .common import get_activation
|
|
|
|
|
|
__all__ = ["CSPResNet"]
|
|
|
|
|
|
|
|
|
donwload_url = {
|
|
|
"s": "https://github.com/lyuwenyu/storage/releases/download/v0.1/CSPResNetb_s_pretrained_from_paddle.pth",
|
|
|
"m": "https://github.com/lyuwenyu/storage/releases/download/v0.1/CSPResNetb_m_pretrained_from_paddle.pth",
|
|
|
"l": "https://github.com/lyuwenyu/storage/releases/download/v0.1/CSPResNetb_l_pretrained_from_paddle.pth",
|
|
|
"x": "https://github.com/lyuwenyu/storage/releases/download/v0.1/CSPResNetb_x_pretrained_from_paddle.pth",
|
|
|
}
|
|
|
|
|
|
|
|
|
class ConvBNLayer(nn.Module):
|
|
|
def __init__(self, ch_in, ch_out, filter_size=3, stride=1, groups=1, padding=0, act=None):
|
|
|
super().__init__()
|
|
|
self.conv = nn.Conv2d(
|
|
|
ch_in, ch_out, filter_size, stride, padding, groups=groups, bias=False
|
|
|
)
|
|
|
self.bn = nn.BatchNorm2d(ch_out)
|
|
|
self.act = get_activation(act)
|
|
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
|
x = self.conv(x)
|
|
|
x = self.bn(x)
|
|
|
x = self.act(x)
|
|
|
return x
|
|
|
|
|
|
|
|
|
class RepVggBlock(nn.Module):
|
|
|
def __init__(self, ch_in, ch_out, act="relu", alpha: bool = False):
|
|
|
super().__init__()
|
|
|
self.ch_in = ch_in
|
|
|
self.ch_out = ch_out
|
|
|
self.conv1 = ConvBNLayer(ch_in, ch_out, 3, stride=1, padding=1, act=None)
|
|
|
self.conv2 = ConvBNLayer(ch_in, ch_out, 1, stride=1, padding=0, act=None)
|
|
|
self.act = get_activation(act)
|
|
|
|
|
|
if alpha:
|
|
|
self.alpha = nn.Parameter(
|
|
|
torch.ones(
|
|
|
1,
|
|
|
)
|
|
|
)
|
|
|
else:
|
|
|
self.alpha = None
|
|
|
|
|
|
def forward(self, x):
|
|
|
if hasattr(self, "conv"):
|
|
|
y = self.conv(x)
|
|
|
else:
|
|
|
if self.alpha:
|
|
|
y = self.conv1(x) + self.alpha * self.conv2(x)
|
|
|
else:
|
|
|
y = self.conv1(x) + self.conv2(x)
|
|
|
y = self.act(y)
|
|
|
return y
|
|
|
|
|
|
def convert_to_deploy(self):
|
|
|
if not hasattr(self, "conv"):
|
|
|
self.conv = nn.Conv2d(self.ch_in, self.ch_out, 3, 1, padding=1)
|
|
|
|
|
|
kernel, bias = self.get_equivalent_kernel_bias()
|
|
|
self.conv.weight.data = kernel
|
|
|
self.conv.bias.data = bias
|
|
|
|
|
|
def get_equivalent_kernel_bias(self):
|
|
|
kernel3x3, bias3x3 = self._fuse_bn_tensor(self.conv1)
|
|
|
kernel1x1, bias1x1 = self._fuse_bn_tensor(self.conv2)
|
|
|
|
|
|
if self.alpha:
|
|
|
return kernel3x3 + self.alpha * self._pad_1x1_to_3x3_tensor(
|
|
|
kernel1x1
|
|
|
), bias3x3 + self.alpha * bias1x1
|
|
|
else:
|
|
|
return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1), bias3x3 + bias1x1
|
|
|
|
|
|
def _pad_1x1_to_3x3_tensor(self, kernel1x1):
|
|
|
if kernel1x1 is None:
|
|
|
return 0
|
|
|
else:
|
|
|
return F.pad(kernel1x1, [1, 1, 1, 1])
|
|
|
|
|
|
def _fuse_bn_tensor(self, branch: ConvBNLayer):
|
|
|
if branch is None:
|
|
|
return 0, 0
|
|
|
kernel = branch.conv.weight
|
|
|
running_mean = branch.norm.running_mean
|
|
|
running_var = branch.norm.running_var
|
|
|
gamma = branch.norm.weight
|
|
|
beta = branch.norm.bias
|
|
|
eps = branch.norm.eps
|
|
|
std = (running_var + eps).sqrt()
|
|
|
t = (gamma / std).reshape(-1, 1, 1, 1)
|
|
|
return kernel * t, beta - running_mean * gamma / std
|
|
|
|
|
|
|
|
|
class BasicBlock(nn.Module):
|
|
|
def __init__(self, ch_in, ch_out, act="relu", shortcut=True, use_alpha=False):
|
|
|
super().__init__()
|
|
|
assert ch_in == ch_out
|
|
|
self.conv1 = ConvBNLayer(ch_in, ch_out, 3, stride=1, padding=1, act=act)
|
|
|
self.conv2 = RepVggBlock(ch_out, ch_out, act=act, alpha=use_alpha)
|
|
|
self.shortcut = shortcut
|
|
|
|
|
|
def forward(self, x):
|
|
|
y = self.conv1(x)
|
|
|
y = self.conv2(y)
|
|
|
if self.shortcut:
|
|
|
return x + y
|
|
|
else:
|
|
|
return y
|
|
|
|
|
|
|
|
|
class EffectiveSELayer(nn.Module):
|
|
|
"""Effective Squeeze-Excitation
|
|
|
From `CenterMask : Real-Time Anchor-Free Instance Segmentation` - https://arxiv.org/abs/1911.06667
|
|
|
"""
|
|
|
|
|
|
def __init__(self, channels, act="hardsigmoid"):
|
|
|
super(EffectiveSELayer, self).__init__()
|
|
|
self.fc = nn.Conv2d(channels, channels, kernel_size=1, padding=0)
|
|
|
self.act = get_activation(act)
|
|
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
|
x_se = x.mean((2, 3), keepdim=True)
|
|
|
x_se = self.fc(x_se)
|
|
|
x_se = self.act(x_se)
|
|
|
return x * x_se
|
|
|
|
|
|
|
|
|
class CSPResStage(nn.Module):
|
|
|
def __init__(self, block_fn, ch_in, ch_out, n, stride, act="relu", attn="eca", use_alpha=False):
|
|
|
super().__init__()
|
|
|
ch_mid = (ch_in + ch_out) // 2
|
|
|
if stride == 2:
|
|
|
self.conv_down = ConvBNLayer(ch_in, ch_mid, 3, stride=2, padding=1, act=act)
|
|
|
else:
|
|
|
self.conv_down = None
|
|
|
self.conv1 = ConvBNLayer(ch_mid, ch_mid // 2, 1, act=act)
|
|
|
self.conv2 = ConvBNLayer(ch_mid, ch_mid // 2, 1, act=act)
|
|
|
self.blocks = nn.Sequential(
|
|
|
*[
|
|
|
block_fn(ch_mid // 2, ch_mid // 2, act=act, shortcut=True, use_alpha=use_alpha)
|
|
|
for i in range(n)
|
|
|
]
|
|
|
)
|
|
|
if attn:
|
|
|
self.attn = EffectiveSELayer(ch_mid, act="hardsigmoid")
|
|
|
else:
|
|
|
self.attn = None
|
|
|
|
|
|
self.conv3 = ConvBNLayer(ch_mid, ch_out, 1, act=act)
|
|
|
|
|
|
def forward(self, x):
|
|
|
if self.conv_down is not None:
|
|
|
x = self.conv_down(x)
|
|
|
y1 = self.conv1(x)
|
|
|
y2 = self.blocks(self.conv2(x))
|
|
|
y = torch.concat([y1, y2], dim=1)
|
|
|
if self.attn is not None:
|
|
|
y = self.attn(y)
|
|
|
y = self.conv3(y)
|
|
|
return y
|
|
|
|
|
|
|
|
|
@register()
|
|
|
class CSPResNet(nn.Module):
|
|
|
layers = [3, 6, 6, 3]
|
|
|
channels = [64, 128, 256, 512, 1024]
|
|
|
model_cfg = {
|
|
|
"s": {
|
|
|
"depth_mult": 0.33,
|
|
|
"width_mult": 0.50,
|
|
|
},
|
|
|
"m": {
|
|
|
"depth_mult": 0.67,
|
|
|
"width_mult": 0.75,
|
|
|
},
|
|
|
"l": {
|
|
|
"depth_mult": 1.00,
|
|
|
"width_mult": 1.00,
|
|
|
},
|
|
|
"x": {
|
|
|
"depth_mult": 1.33,
|
|
|
"width_mult": 1.25,
|
|
|
},
|
|
|
}
|
|
|
|
|
|
def __init__(
|
|
|
self,
|
|
|
name: str,
|
|
|
act="silu",
|
|
|
return_idx=[1, 2, 3],
|
|
|
use_large_stem=True,
|
|
|
use_alpha=False,
|
|
|
pretrained=False,
|
|
|
):
|
|
|
super().__init__()
|
|
|
depth_mult = self.model_cfg[name]["depth_mult"]
|
|
|
width_mult = self.model_cfg[name]["width_mult"]
|
|
|
|
|
|
channels = [max(round(c * width_mult), 1) for c in self.channels]
|
|
|
layers = [max(round(l * depth_mult), 1) for l in self.layers]
|
|
|
act = get_activation(act)
|
|
|
|
|
|
if use_large_stem:
|
|
|
self.stem = nn.Sequential(
|
|
|
OrderedDict(
|
|
|
[
|
|
|
(
|
|
|
"conv1",
|
|
|
ConvBNLayer(3, channels[0] // 2, 3, stride=2, padding=1, act=act),
|
|
|
),
|
|
|
(
|
|
|
"conv2",
|
|
|
ConvBNLayer(
|
|
|
channels[0] // 2, channels[0] // 2, 3, stride=1, padding=1, act=act
|
|
|
),
|
|
|
),
|
|
|
(
|
|
|
"conv3",
|
|
|
ConvBNLayer(
|
|
|
channels[0] // 2, channels[0], 3, stride=1, padding=1, act=act
|
|
|
),
|
|
|
),
|
|
|
]
|
|
|
)
|
|
|
)
|
|
|
else:
|
|
|
self.stem = nn.Sequential(
|
|
|
OrderedDict(
|
|
|
[
|
|
|
(
|
|
|
"conv1",
|
|
|
ConvBNLayer(3, channels[0] // 2, 3, stride=2, padding=1, act=act),
|
|
|
),
|
|
|
(
|
|
|
"conv2",
|
|
|
ConvBNLayer(
|
|
|
channels[0] // 2, channels[0], 3, stride=1, padding=1, act=act
|
|
|
),
|
|
|
),
|
|
|
]
|
|
|
)
|
|
|
)
|
|
|
|
|
|
n = len(channels) - 1
|
|
|
self.stages = nn.Sequential(
|
|
|
OrderedDict(
|
|
|
[
|
|
|
(
|
|
|
str(i),
|
|
|
CSPResStage(
|
|
|
BasicBlock,
|
|
|
channels[i],
|
|
|
channels[i + 1],
|
|
|
layers[i],
|
|
|
2,
|
|
|
act=act,
|
|
|
use_alpha=use_alpha,
|
|
|
),
|
|
|
)
|
|
|
for i in range(n)
|
|
|
]
|
|
|
)
|
|
|
)
|
|
|
|
|
|
self._out_channels = channels[1:]
|
|
|
self._out_strides = [4 * 2**i for i in range(n)]
|
|
|
self.return_idx = return_idx
|
|
|
|
|
|
if pretrained:
|
|
|
if isinstance(pretrained, bool) or "http" in pretrained:
|
|
|
state = torch.hub.load_state_dict_from_url(donwload_url[name], map_location="cpu")
|
|
|
else:
|
|
|
state = torch.load(pretrained, map_location="cpu")
|
|
|
self.load_state_dict(state)
|
|
|
print(f"Load CSPResNet_{name} state_dict")
|
|
|
|
|
|
def forward(self, x):
|
|
|
x = self.stem(x)
|
|
|
outs = []
|
|
|
for idx, stage in enumerate(self.stages):
|
|
|
x = stage(x)
|
|
|
if idx in self.return_idx:
|
|
|
outs.append(x)
|
|
|
|
|
|
return outs
|
|
|
|
