Add verbose option to pytorch hub models

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

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

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):