Param count

models.py

@@ -31,6 +31,9 @@ def getComp(self,W,H):
 
         return self.size*self.size*W*H*self.inch*self.ch*2 + (W*H*self.ch*4 if self.doBN else 0), W, H
 
+    def getParams(self):
+        return self.ch*(self.inch*self.size*self.size + 4 if self.doBN else 1)
+
 class YOLOLayer(nn.Module):
     """Detection layer"""
 
@@ -292,3 +295,9 @@ def get_computations(self,pruned = False):
 
         return computations
 
+    def getParams(self):
+        params = sum([layer.getParams() for layer in self.downPart if layer is not None])
+        params += 64*10*2 * (2 if self.bn else 1)
+        params += 128*10//2 * (2 if self.bn else 1)
+        return params
+

test.py

@@ -154,6 +154,8 @@ def getmAP(all_annotations,all_detections):
     for path in weights_path:
         print(path)
         model = ROBO(inch=channels,bn=opt.bn, halfRes=opt.hr)
+        print(model.getParams())
+        exit(0)
         model.load_state_dict(torch.load(path,map_location={'cuda:0': 'cpu'}))
 
         print(count_zero_weights(model))