改进yolov8|FasterNet更换主干网络，跑得飞快！！ - ToB企服应用市场:ToB评测及商务社交产业平台

# --------------------------FasterNet----------------------------
from timm.models.layers import DropPath
class Partial_conv3(nn.Module):
def __init__(self, dim, n_div, forward):
super().__init__()
self.dim_conv3 = dim // n_div
self.dim_untouched = dim - self.dim_conv3
self.partial_conv3 = nn.Conv2d(self.dim_conv3, self.dim_conv3, 3, 1, 1, bias=False)
if forward == 'slicing':
self.forward = self.forward_slicing
elif forward == 'split_cat':
self.forward = self.forward_split_cat
else:
raise NotImplementedError
def forward_slicing(self, x):
# only for inference
x = x.clone() # !!! Keep the original input intact for the residual connection later
x[:, :self.dim_conv3, :, :] = self.partial_conv3(x[:, :self.dim_conv3, :, :])
return x
def forward_split_cat(self, x):
# for training/inference
x1, x2 = torch.split(x, [self.dim_conv3, self.dim_untouched], dim=1)
x1 = self.partial_conv3(x1)
x = torch.cat((x1, x2), 1)
return x
class MLPBlock(nn.Module):
def __init__(self,
dim,
n_div,
mlp_ratio,
drop_path,
layer_scale_init_value,
act_layer,
norm_layer,
pconv_fw_type
):
super().__init__()
self.dim = dim
self.mlp_ratio = mlp_ratio
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.n_div = n_div
mlp_hidden_dim = int(dim * mlp_ratio)
mlp_layer = [
nn.Conv2d(dim, mlp_hidden_dim, 1, bias=False),
norm_layer(mlp_hidden_dim),
act_layer(),
nn.Conv2d(mlp_hidden_dim, dim, 1, bias=False)
]
self.mlp = nn.Sequential(*mlp_layer)
self.spatial_mixing = Partial_conv3(
dim,
n_div,
pconv_fw_type
)
if layer_scale_init_value > 0:
self.layer_scale = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
self.forward = self.forward_layer_scale
else:
self.forward = self.forward
def forward(self, x):
shortcut = x
x = self.spatial_mixing(x)
x = shortcut + self.drop_path(self.mlp(x))
return x
def forward_layer_scale(self, x):
shortcut = x
x = self.spatial_mixing(x)
x = shortcut + self.drop_path(
self.layer_scale.unsqueeze(-1).unsqueeze(-1) * self.mlp(x))
return x
class BasicStage(nn.Module):
def __init__(self,
dim,
depth=1,
n_div=4,
mlp_ratio=2,
layer_scale_init_value=0,
norm_layer=nn.BatchNorm2d,
act_layer=nn.ReLU,
pconv_fw_type='split_cat'
):
super().__init__()
dpr = [x.item()
for x in torch.linspace(0, 0.0, sum([1, 2, 8, 2]))]
blocks_list = [
MLPBlock(
dim=dim,
n_div=n_div,
mlp_ratio=mlp_ratio,
drop_path=dpr[i],
layer_scale_init_value=layer_scale_init_value,
norm_layer=norm_layer,
act_layer=act_layer,
pconv_fw_type=pconv_fw_type
)
for i in range(depth)
]
self.blocks = nn.Sequential(*blocks_list)
def forward(self, x):
x = self.blocks(x)
return x
class PatchEmbed_FasterNet(nn.Module):
def __init__(self, in_chans, embed_dim, patch_size, patch_stride, norm_layer=nn.BatchNorm2d):
super().__init__()
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_stride, bias=False)
if norm_layer is not None:
self.norm = norm_layer(embed_dim)
else:
self.norm = nn.Identity()
def forward(self, x):
x = self.norm(self.proj(x))
return x
def fuseforward(self, x):
x = self.proj(x)
return x
class PatchMerging_FasterNet(nn.Module):
def __init__(self, dim, out_dim, k, patch_stride2, norm_layer=nn.BatchNorm2d):
super().__init__()
self.reduction = nn.Conv2d(dim, out_dim, kernel_size=k, stride=patch_stride2, bias=False)
if norm_layer is not None:
self.norm = norm_layer(out_dim)
else:
self.norm = nn.Identity()
def forward(self, x):
x = self.norm(self.reduction(x))
return x
def fuseforward(self, x):
x = self.reduction(x)
return x

复制代码

from .block import (....,BasicStage,PatchEmbed_FasterNet,PatchMerging_FasterNet)

复制代码

from ultralytics.nn.modules import (....BasicStage, PatchEmbed_FasterNet, PatchMerging_FasterNet)

复制代码

if m in (... BasicStage, PatchEmbed_FasterNet, PatchMerging_FasterNet):

复制代码

elif m in [BasicStage]:
args.pop(1)

复制代码

if type(m) is PatchEmbed_FasterNet:
m.proj = fuse_conv_and_bn(m.proj, m.norm)
delattr(m, 'norm') # remove BN
m.forward = m.fuseforward
if type(m) is PatchMerging_FasterNet:
m.reduction = fuse_conv_and_bn(m.reduction, m.norm)
delattr(m, 'norm') # remove BN
m.forward = m.fuseforward

复制代码