diff --git a/hubconf.py b/hubconf.py index d89502f4ee76..e42d0b59bd2a 100644 --- a/hubconf.py +++ b/hubconf.py @@ -16,10 +16,9 @@ dependencies = ['torch', 'yaml'] check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('pycocotools', 'thop')) -set_logging() -def create(name, pretrained, channels, classes, autoshape): +def create(name, pretrained, channels, classes, autoshape, verbose): """Creates a specified YOLOv5 model Arguments: @@ -32,6 +31,8 @@ def create(name, pretrained, channels, classes, autoshape): pytorch model """ try: + set_logging(verbose=verbose) + cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path model = Model(cfg, channels, classes) if pretrained: @@ -55,7 +56,7 @@ def create(name, pretrained, channels, classes, autoshape): raise Exception(s) from e -def custom(path_or_model='path/to/model.pt', autoshape=True): +def custom(path_or_model='path/to/model.pt', autoshape=True, verbose=True): """YOLOv5-custom model https://github.com/ultralytics/yolov5 Arguments (3 options): @@ -66,6 +67,8 @@ def custom(path_or_model='path/to/model.pt', autoshape=True): Returns: pytorch model """ + set_logging(verbose=verbose) + model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint if isinstance(model, dict): model = model['ema' if model.get('ema') else 'model'] # load model @@ -79,49 +82,49 @@ def custom(path_or_model='path/to/model.pt', autoshape=True): return hub_model.to(device) -def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True): +def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): # YOLOv5-small model https://github.com/ultralytics/yolov5 - return create('yolov5s', pretrained, channels, classes, autoshape) + return create('yolov5s', pretrained, channels, classes, autoshape, verbose) -def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True): +def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): # YOLOv5-medium model https://github.com/ultralytics/yolov5 - return create('yolov5m', pretrained, channels, classes, autoshape) + return create('yolov5m', pretrained, channels, classes, autoshape, verbose) -def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True): +def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): # YOLOv5-large model https://github.com/ultralytics/yolov5 - return create('yolov5l', pretrained, channels, classes, autoshape) + return create('yolov5l', pretrained, channels, classes, autoshape, verbose) -def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True): +def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): # YOLOv5-xlarge model https://github.com/ultralytics/yolov5 - return create('yolov5x', pretrained, channels, classes, autoshape) + return create('yolov5x', pretrained, channels, classes, autoshape, verbose) -def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True): - # YOLOv5-small model https://github.com/ultralytics/yolov5 - return create('yolov5s6', pretrained, channels, classes, autoshape) +def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): + # YOLOv5-small-P6 model https://github.com/ultralytics/yolov5 + return create('yolov5s6', pretrained, channels, classes, autoshape, verbose) -def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True): - # YOLOv5-medium model https://github.com/ultralytics/yolov5 - return create('yolov5m6', pretrained, channels, classes, autoshape) +def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): + # YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5 + return create('yolov5m6', pretrained, channels, classes, autoshape, verbose) -def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True): - # YOLOv5-large model https://github.com/ultralytics/yolov5 - return create('yolov5l6', pretrained, channels, classes, autoshape) +def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): + # YOLOv5-large-P6 model https://github.com/ultralytics/yolov5 + return create('yolov5l6', pretrained, channels, classes, autoshape, verbose) -def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True): - # YOLOv5-xlarge model https://github.com/ultralytics/yolov5 - return create('yolov5x6', pretrained, channels, classes, autoshape) +def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): + # YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5 + return create('yolov5x6', pretrained, channels, classes, autoshape, verbose) if __name__ == '__main__': - model = create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True) # pretrained example - # model = custom(path_or_model='path/to/model.pt') # custom example + model = create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained + # model = custom(path_or_model='path/to/model.pt') # custom # Verify inference import cv2 diff --git a/models/yolo.py b/models/yolo.py index 36fa27e89134..dd505e22a68d 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -84,7 +84,7 @@ def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, i self.yaml['anchors'] = round(anchors) # override yaml value self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist self.names = [str(i) for i in range(self.yaml['nc'])] # default names - # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) + # logger.info([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) # Build strides, anchors m = self.model[-1] # Detect() @@ -95,7 +95,7 @@ def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, i check_anchor_order(m) self.stride = m.stride self._initialize_biases() # only run once - # print('Strides: %s' % m.stride.tolist()) + # logger.info('Strides: %s' % m.stride.tolist()) # Init weights, biases initialize_weights(self) @@ -134,13 +134,13 @@ def forward_once(self, x, profile=False): for _ in range(10): _ = m(x) dt.append((time_synchronized() - t) * 100) - print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type)) + logger.info('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type)) x = m(x) # run y.append(x if m.i in self.save else None) # save output if profile: - print('%.1fms total' % sum(dt)) + logger.info('%.1fms total' % sum(dt)) return x def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency @@ -157,15 +157,15 @@ def _print_biases(self): m = self.model[-1] # Detect() module for mi in m.m: # from b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) - print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) + logger.info(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) # def _print_weights(self): # for m in self.model.modules(): # if type(m) is Bottleneck: - # print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights + # logger.info('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers - print('Fusing layers... ') + logger.info('Fusing layers... ') for m in self.model.modules(): if type(m) is Conv and hasattr(m, 'bn'): m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv @@ -177,19 +177,19 @@ def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers def nms(self, mode=True): # add or remove NMS module present = type(self.model[-1]) is NMS # last layer is NMS if mode and not present: - print('Adding NMS... ') + logger.info('Adding NMS... ') 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 self.eval() elif not mode and present: - print('Removing NMS... ') + logger.info('Removing NMS... ') self.model = self.model[:-1] # remove return self def autoshape(self): # add autoShape module - print('Adding autoShape... ') + logger.info('Adding autoShape... ') m = autoShape(self) # wrap model copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes return m @@ -272,6 +272,6 @@ def parse_model(d, ch): # model_dict, input_channels(3) # Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898) # from torch.utils.tensorboard import SummaryWriter # tb_writer = SummaryWriter('.') - # print("Run 'tensorboard --logdir=models' to view tensorboard at http://localhost:6006/") + # logger.info("Run 'tensorboard --logdir=models' to view tensorboard at http://localhost:6006/") # tb_writer.add_graph(torch.jit.trace(model, img, strict=False), []) # add model graph # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard diff --git a/utils/general.py b/utils/general.py index 8efeb5ea59cc..f77ae3331538 100755 --- a/utils/general.py +++ b/utils/general.py @@ -32,10 +32,10 @@ os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads -def set_logging(rank=-1): +def set_logging(rank=-1, verbose=True): logging.basicConfig( format="%(message)s", - level=logging.INFO if rank in [-1, 0] else logging.WARN) + level=logging.INFO if (verbose and rank in [-1, 0]) else logging.WARN) def init_seeds(seed=0):