model.add_nms() method

hubconf.py

@@ -37,10 +37,7 @@ def create(name, pretrained, channels, classes):
             state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape}  # filter
             model.load_state_dict(state_dict, strict=False)  # load
 
-            m = NMS()
-            m.f = -1  # from
-            m.i = model.model[-1].i + 1  # index
-            model.model.add_module(name='%s' % m.i, module=m)  # add NMS
+            model.add_nms()  # add NMS module
             model.eval()
         return model
 

models/yolo.py

@@ -7,7 +7,7 @@ from pathlib import Path
 import torch
 import torch.nn as nn
 
-from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat
+from models.common import Conv, Bottleneck, SPP, DWConv, Focus, BottleneckCSP, Concat, NMS
 from models.experimental import MixConv2d, CrossConv, C3
 from utils.general import check_anchor_order, make_divisible, check_file, set_logging
 from utils.torch_utils import (
@@ -168,6 +168,15 @@ class Model(nn.Module):
         self.info()
         return self
 
+    def add_nms(self):  # fuse model Conv2d() + BatchNorm2d() layers
+        if type(self.model[-1]) is not NMS:  # if missing NMS
+            print('Adding NMS module... ')
+            m = NMS()  # module
+            m.f = -1  # from
+            m.i = self.model[-1].i + 1  # index
+            self.model.add_module(name='%s' % m.i, module=m)  # add
+        return self
+
     def info(self, verbose=False):  # print model information
         model_info(self, verbose)