lucytuan commited on
Commit
9a3d99f
·
1 Parent(s): b67aac7

✨ [Add] RepConv module in module.py

Browse files
Files changed (1) hide show
  1. 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
- # https://github.com/DingXiaoH/RepVGG
 
105
  def __init__(
106
- self, in_channels, out_channels, kernel_size=3, padding=None, stride=1, groups=1, act=nn.SiLU(), deploy=False
 
 
 
 
 
 
107
  ):
108
-
109
  super().__init__()
110
- self.deploy = deploy
111
- self.conv1 = Conv(in_channels, out_channels, kernel_size, stride=stride, groups=groups, activation=False)
112
- self.conv2 = Conv(in_channels, out_channels, 1, stride=stride, groups=groups, activation=False)
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):