diff --git a/hubconf.py b/hubconf.py index ba150e0..168b405 100644 --- a/hubconf.py +++ b/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 diff --git a/models/yolo.py b/models/yolo.py index 4b2606d..348c418 100644 --- a/models/yolo.py +++ b/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)