create code

nets/CSPdarknet.py

@@ -15,9 +15,9 @@ class SiLU(nn.Module):
 class Conv(nn.Module):
     def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=SiLU()):  # ch_in, ch_out, kernel, stride, padding, groups
         super(Conv, self).__init__()
-        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
-        self.bn = nn.BatchNorm2d(c2)
-        self.act = nn.LeakyReLU(0.1, inplace=True) if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
+        self.conv   = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
+        self.bn     = nn.BatchNorm2d(c2, eps=0.001, momentum=0.03)
+        self.act    = nn.LeakyReLU(0.1, inplace=True) if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
 
     def forward(self, x):
         return self.act(self.bn(self.conv(x)))
@@ -116,10 +116,7 @@ class CSPDarknet(nn.Module):
             print("Load weights from ", url.split('/')[-1])
 
     def forward(self, x):
-        x = self.stem[0](x)
-        
-        x = self.stem[1](x)
-        x = self.stem[2](x)
+        x = self.stem(x)
         x = self.dark2(x)
         #-----------------------------------------------#
         #   dark3的输出为80, 80, 256，是一个有效特征层

nets/yolo.py

@@ -45,14 +45,14 @@ class RepConv(nn.Module):
         if deploy:
             self.rbr_reparam    = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=True)
         else:
-            self.rbr_identity   = (nn.BatchNorm2d(num_features=c1) if c2 == c1 and s == 1 else None)
+            self.rbr_identity   = (nn.BatchNorm2d(num_features=c1, eps=0.001, momentum=0.03) if c2 == c1 and s == 1 else None)
             self.rbr_dense      = nn.Sequential(
                 nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False),
-                nn.BatchNorm2d(num_features=c2),
+                nn.BatchNorm2d(num_features=c2, eps=0.001, momentum=0.03),
             )
             self.rbr_1x1        = nn.Sequential(
                 nn.Conv2d( c1, c2, 1, s, padding_11, groups=g, bias=False),
-                nn.BatchNorm2d(num_features=c2),
+                nn.BatchNorm2d(num_features=c2, eps=0.001, momentum=0.03),
             )
 
     def forward(self, inputs):

yolo.py

@@ -100,7 +100,7 @@ class YOLO(object):
         self.net    = YoloBody(self.anchors_mask, self.num_classes)
         device      = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.net.load_state_dict(torch.load(self.model_path, map_location=device))
-        self.net    = self.net.eval()
+        self.net    = self.net.fuse().eval()
         print('{} model, and classes loaded.'.format(self.model_path))
         if not onnx:
             if self.cuda: