✨ [Add] RepConv module in module.py
Browse files- yolo/model/module.py +13 -38
yolo/model/module.py
CHANGED
@@ -101,50 +101,25 @@ class SPPELAN(nn.Module):
|
|
101 |
|
102 |
# RepVGG
|
103 |
class RepConv(nn.Module):
|
104 |
-
|
|
|
105 |
def __init__(
|
106 |
-
self,
|
|
|
|
|
|
|
|
|
|
|
|
|
107 |
):
|
108 |
-
|
109 |
super().__init__()
|
110 |
-
self.
|
111 |
-
self.conv1 = Conv(in_channels, out_channels, kernel_size,
|
112 |
-
self.conv2 = Conv(in_channels, out_channels, 1,
|
113 |
-
self.act = act if isinstance(act, nn.Module) else nn.Identity()
|
114 |
|
115 |
-
def forward(self, x):
|
116 |
return self.act(self.conv1(x) + self.conv2(x))
|
117 |
|
118 |
-
def forward_fuse(self, x):
|
119 |
-
return self.act(self.conv(x))
|
120 |
-
|
121 |
-
# to be implement
|
122 |
-
# def fuse_convs(self):
|
123 |
-
def fuse_conv_bn(self, conv, bn):
|
124 |
-
|
125 |
-
std = (bn.running_var + bn.eps).sqrt()
|
126 |
-
bias = bn.bias - bn.running_mean * bn.weight / std
|
127 |
-
|
128 |
-
t = (bn.weight / std).reshape(-1, 1, 1, 1)
|
129 |
-
weights = conv.weight * t
|
130 |
-
|
131 |
-
bn = nn.Identity()
|
132 |
-
conv = nn.Conv2d(
|
133 |
-
in_channels=conv.in_channels,
|
134 |
-
out_channels=conv.out_channels,
|
135 |
-
kernel_size=conv.kernel_size,
|
136 |
-
stride=conv.stride,
|
137 |
-
padding=conv.padding,
|
138 |
-
dilation=conv.dilation,
|
139 |
-
groups=conv.groups,
|
140 |
-
bias=True,
|
141 |
-
padding_mode=conv.padding_mode,
|
142 |
-
)
|
143 |
-
|
144 |
-
conv.weight = torch.nn.Parameter(weights)
|
145 |
-
conv.bias = torch.nn.Parameter(bias)
|
146 |
-
return conv
|
147 |
-
|
148 |
|
149 |
# ResNet
|
150 |
class Res(nn.Module):
|
|
|
101 |
|
102 |
# RepVGG
|
103 |
class RepConv(nn.Module):
|
104 |
+
"""A convolutional block that combines two convolution layers (kernel and point-wise)."""
|
105 |
+
|
106 |
def __init__(
|
107 |
+
self,
|
108 |
+
in_channels: int,
|
109 |
+
out_channels: int,
|
110 |
+
kernel_size: _size_2_t = 3,
|
111 |
+
*,
|
112 |
+
activation: Optional[str] = "SiLU",
|
113 |
+
**kwargs
|
114 |
):
|
|
|
115 |
super().__init__()
|
116 |
+
self.act = get_activation(activation)
|
117 |
+
self.conv1 = Conv(in_channels, out_channels, kernel_size, activation=False, **kwargs)
|
118 |
+
self.conv2 = Conv(in_channels, out_channels, 1, activation=False, **kwargs)
|
|
|
119 |
|
120 |
+
def forward(self, x: Tensor) -> Tensor:
|
121 |
return self.act(self.conv1(x) + self.conv2(x))
|
122 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
123 |
|
124 |
# ResNet
|
125 |
class Res(nn.Module):
|