Update get_loggers() (ultralytics#4854)

* Update `set_logging()` * Update export.py * pre-commit fixes * Update LoadImages * Update LoadStreams * Update print_args * Single LOGGER definition * yolo.py fix Co-authored-by: pre-commit <pre-commit@example.com>
fhkiel-mlaip · Nov 1, 2021 · e6bcd7a · e6bcd7a
1 parent 2bb3ac0
commit e6bcd7a
Show file tree

Hide file tree

Showing 8 changed files with 84 additions and 91 deletions.
diff --git a/detect.py b/detect.py
@@ -25,8 +25,7 @@
 from models.experimental import attempt_load
 from utils.datasets import LoadImages, LoadStreams
 from utils.general import apply_classifier, check_img_size, check_imshow, check_requirements, check_suffix, colorstr, \
-    increment_path, non_max_suppression, print_args, save_one_box, scale_coords, set_logging, \
-    strip_optimizer, xyxy2xywh
+    increment_path, non_max_suppression, print_args, save_one_box, scale_coords, strip_optimizer, xyxy2xywh, LOGGER
 from utils.plots import Annotator, colors
 from utils.torch_utils import load_classifier, select_device, time_sync
 
@@ -68,7 +67,6 @@ def run(weights=ROOT / 'yolov5s.pt',  # model.pt path(s)
     (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir
 
     # Initialize
-    set_logging()
     device = select_device(device)
     half &= device.type != 'cpu'  # half precision only supported on CUDA
 
@@ -132,7 +130,7 @@ def wrap_frozen_graph(gd, inputs, outputs):
     if pt and device.type != 'cpu':
         model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.parameters())))  # run once
     dt, seen = [0.0, 0.0, 0.0], 0
-    for path, img, im0s, vid_cap in dataset:
+    for path, img, im0s, vid_cap, s in dataset:
         t1 = time_sync()
         if onnx:
             img = img.astype('float32')
@@ -191,9 +189,10 @@ def wrap_frozen_graph(gd, inputs, outputs):
         for i, det in enumerate(pred):  # per image
             seen += 1
             if webcam:  # batch_size >= 1
-                p, s, im0, frame = path[i], f'{i}: ', im0s[i].copy(), dataset.count
+                p, im0, frame = path[i], im0s[i].copy(), dataset.count
+                s += f'{i}: '
             else:
-                p, s, im0, frame = path, '', im0s.copy(), getattr(dataset, 'frame', 0)
+                p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
 
             p = Path(p)  # to Path
             save_path = str(save_dir / p.name)  # img.jpg
@@ -227,7 +226,7 @@ def wrap_frozen_graph(gd, inputs, outputs):
                             save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
 
             # Print time (inference-only)
-            print(f'{s}Done. ({t3 - t2:.3f}s)')
+            LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
 
             # Stream results
             im0 = annotator.result()
@@ -256,10 +255,10 @@ def wrap_frozen_graph(gd, inputs, outputs):
 
     # Print results
     t = tuple(x / seen * 1E3 for x in dt)  # speeds per image
-    print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
+    LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
     if save_txt or save_img:
         s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
-        print(f"Results saved to {colorstr('bold', save_dir)}{s}")
+        LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
     if update:
         strip_optimizer(weights)  # update model (to fix SourceChangeWarning)
 

diff --git a/export.py b/export.py
@@ -42,23 +42,23 @@
 from models.yolo import Detect
 from utils.activations import SiLU
 from utils.datasets import LoadImages
-from utils.general import colorstr, check_dataset, check_img_size, check_requirements, file_size, print_args, \
-    set_logging, url2file
+from utils.general import check_dataset, check_img_size, check_requirements, colorstr, file_size, print_args, \
+    url2file, LOGGER
 from utils.torch_utils import select_device
 
 
 def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
     # YOLOv5 TorchScript model export
     try:
-        print(f'\n{prefix} starting export with torch {torch.__version__}...')
+        LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
         f = file.with_suffix('.torchscript.pt')
 
         ts = torch.jit.trace(model, im, strict=False)
         (optimize_for_mobile(ts) if optimize else ts).save(f)
 
-        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
     except Exception as e:
-        print(f'{prefix} export failure: {e}')
+        LOGGER.info(f'{prefix} export failure: {e}')
 
 
 def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
@@ -67,7 +67,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
         check_requirements(('onnx',))
         import onnx
 
-        print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
+        LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
         f = file.with_suffix('.onnx')
 
         torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
@@ -82,27 +82,27 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst
         # Checks
         model_onnx = onnx.load(f)  # load onnx model
         onnx.checker.check_model(model_onnx)  # check onnx model
-        # print(onnx.helper.printable_graph(model_onnx.graph))  # print
+        # LOGGER.info(onnx.helper.printable_graph(model_onnx.graph))  # print
 
         # Simplify
         if simplify:
             try:
                 check_requirements(('onnx-simplifier',))
                 import onnxsim
 
-                print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
+                LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
                 model_onnx, check = onnxsim.simplify(
                     model_onnx,
                     dynamic_input_shape=dynamic,
                     input_shapes={'images': list(im.shape)} if dynamic else None)
                 assert check, 'assert check failed'
                 onnx.save(model_onnx, f)
             except Exception as e:
-                print(f'{prefix} simplifier failure: {e}')
-        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
-        print(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'")
+                LOGGER.info(f'{prefix} simplifier failure: {e}')
+        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        LOGGER.info(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'")
     except Exception as e:
-        print(f'{prefix} export failure: {e}')
+        LOGGER.info(f'{prefix} export failure: {e}')
 
 
 def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
@@ -112,17 +112,17 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
         check_requirements(('coremltools',))
         import coremltools as ct
 
-        print(f'\n{prefix} starting export with coremltools {ct.__version__}...')
+        LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
         f = file.with_suffix('.mlmodel')
 
         model.train()  # CoreML exports should be placed in model.train() mode
         ts = torch.jit.trace(model, im, strict=False)  # TorchScript model
         ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255.0, bias=[0, 0, 0])])
         ct_model.save(f)
 
-        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
     except Exception as e:
-        print(f'\n{prefix} export failure: {e}')
+        LOGGER.info(f'\n{prefix} export failure: {e}')
 
     return ct_model
 
@@ -137,7 +137,7 @@ def export_saved_model(model, im, file, dynamic,
         from tensorflow import keras
         from models.tf import TFModel, TFDetect
 
-        print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
+        LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
         f = str(file).replace('.pt', '_saved_model')
         batch_size, ch, *imgsz = list(im.shape)  # BCHW
 
@@ -151,9 +151,9 @@ def export_saved_model(model, im, file, dynamic,
         keras_model.summary()
         keras_model.save(f, save_format='tf')
 
-        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
     except Exception as e:
-        print(f'\n{prefix} export failure: {e}')
+        LOGGER.info(f'\n{prefix} export failure: {e}')
 
     return keras_model
 
@@ -164,7 +164,7 @@ def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
         import tensorflow as tf
         from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
 
-        print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
+        LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
         f = file.with_suffix('.pb')
 
         m = tf.function(lambda x: keras_model(x))  # full model
@@ -173,9 +173,9 @@ def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')):
         frozen_func.graph.as_graph_def()
         tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False)
 
-        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
     except Exception as e:
-        print(f'\n{prefix} export failure: {e}')
+        LOGGER.info(f'\n{prefix} export failure: {e}')
 
 
 def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')):
@@ -184,7 +184,7 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
         import tensorflow as tf
         from models.tf import representative_dataset_gen
 
-        print(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
+        LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...')
         batch_size, ch, *imgsz = list(im.shape)  # BCHW
         f = str(file).replace('.pt', '-fp16.tflite')
 
@@ -204,10 +204,10 @@ def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('Te
 
         tflite_model = converter.convert()
         open(f, "wb").write(tflite_model)
-        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
 
     except Exception as e:
-        print(f'\n{prefix} export failure: {e}')
+        LOGGER.info(f'\n{prefix} export failure: {e}')
 
 
 def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
@@ -217,7 +217,7 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
         import re
         import tensorflowjs as tfjs
 
-        print(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
+        LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...')
         f = str(file).replace('.pt', '_web_model')  # js dir
         f_pb = file.with_suffix('.pb')  # *.pb path
         f_json = f + '/model.json'  # *.json path
@@ -240,9 +240,9 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')):
                 json)
             j.write(subst)
 
-        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
+        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
     except Exception as e:
-        print(f'\n{prefix} export failure: {e}')
+        LOGGER.info(f'\n{prefix} export failure: {e}')
 
 
 @torch.no_grad()
@@ -297,7 +297,7 @@ def run(data=ROOT / 'data/coco128.yaml',  # 'dataset.yaml path'
 
     for _ in range(2):
         y = model(im)  # dry runs
-    print(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")
+    LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")
 
     # Exports
     if 'torchscript' in include:
@@ -322,9 +322,9 @@ def run(data=ROOT / 'data/coco128.yaml',  # 'dataset.yaml path'
             export_tfjs(model, im, file)
 
     # Finish
-    print(f'\nExport complete ({time.time() - t:.2f}s)'
-          f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
-          f'\nVisualize with https://netron.app')
+    LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
+                f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
+                f'\nVisualize with https://netron.app')
 
 
 def parse_opt():
@@ -355,7 +355,6 @@ def parse_opt():
 
 
 def main(opt):
-    set_logging()
     run(**vars(opt))
 
 

diff --git a/models/tf.py b/models/tf.py
@@ -31,11 +31,9 @@
 from models.common import Bottleneck, BottleneckCSP, Concat, Conv, C3, DWConv, Focus, SPP, SPPF, autopad
 from models.experimental import CrossConv, MixConv2d, attempt_load
 from models.yolo import Detect
-from utils.general import make_divisible, print_args, set_logging
+from utils.general import make_divisible, print_args, LOGGER
 from utils.activations import SiLU
 
-LOGGER = logging.getLogger(__name__)
-
 
 class TFBN(keras.layers.Layer):
     # TensorFlow BatchNormalization wrapper
@@ -336,7 +334,7 @@ def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, model=None, imgsz=(640, 64
 
         # Define model
         if nc and nc != self.yaml['nc']:
-            print(f"Overriding {cfg} nc={self.yaml['nc']} with nc={nc}")
+            LOGGER.info(f"Overriding {cfg} nc={self.yaml['nc']} with nc={nc}")
             self.yaml['nc'] = nc  # override yaml value
         self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch], model=model, imgsz=imgsz)
 
@@ -457,7 +455,6 @@ def parse_opt():
 
 
 def main(opt):
-    set_logging()
     run(**vars(opt))
 
 

diff --git a/models/yolo.py b/models/yolo.py
@@ -20,7 +20,7 @@
 from models.common import *
 from models.experimental import *
 from utils.autoanchor import check_anchor_order
-from utils.general import check_yaml, make_divisible, print_args, set_logging, check_version
+from utils.general import check_version, check_yaml, make_divisible, print_args, LOGGER
 from utils.plots import feature_visualization
 from utils.torch_utils import copy_attr, fuse_conv_and_bn, initialize_weights, model_info, scale_img, \
     select_device, time_sync
@@ -30,8 +30,6 @@
 except ImportError:
     thop = None
 
-LOGGER = logging.getLogger(__name__)
-
 
 class Detect(nn.Module):
     stride = None  # strides computed during build
@@ -311,7 +309,6 @@ def parse_model(d, ch):  # model_dict, input_channels(3)
     opt = parser.parse_args()
     opt.cfg = check_yaml(opt.cfg)  # check YAML
     print_args(FILE.stem, opt)
-    set_logging()
     device = select_device(opt.device)
 
     # Create model

diff --git a/train.py b/train.py
@@ -40,7 +40,7 @@
 from utils.datasets import create_dataloader
 from utils.general import labels_to_class_weights, increment_path, labels_to_image_weights, init_seeds, \
     strip_optimizer, get_latest_run, check_dataset, check_git_status, check_img_size, check_requirements, \
-    check_file, check_yaml, check_suffix, print_args, print_mutation, set_logging, one_cycle, colorstr, methods
+    check_file, check_yaml, check_suffix, print_args, print_mutation, one_cycle, colorstr, methods, LOGGER
 from utils.downloads import attempt_download
 from utils.loss import ComputeLoss
 from utils.plots import plot_labels, plot_evolve
@@ -51,7 +51,6 @@
 from utils.loggers import Loggers
 from utils.callbacks import Callbacks
 
-LOGGER = logging.getLogger(__name__)
 LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))  # https://pytorch.org/docs/stable/elastic/run.html
 RANK = int(os.getenv('RANK', -1))
 WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
@@ -129,7 +128,7 @@ def train(hyp,  # path/to/hyp.yaml or hyp dictionary
     for k, v in model.named_parameters():
         v.requires_grad = True  # train all layers
         if any(x in k for x in freeze):
-            print(f'freezing {k}')
+            LOGGER.info(f'freezing {k}')
             v.requires_grad = False
 
     # Image size
@@ -485,7 +484,6 @@ def parse_opt(known=False):
 
 def main(opt, callbacks=Callbacks()):
     # Checks
-    set_logging(RANK)
     if RANK in [-1, 0]:
         print_args(FILE.stem, opt)
         check_git_status()
@@ -609,9 +607,9 @@ def main(opt, callbacks=Callbacks()):
 
         # Plot results
         plot_evolve(evolve_csv)
-        print(f'Hyperparameter evolution finished\n'
-              f"Results saved to {colorstr('bold', save_dir)}\n"
-              f'Use best hyperparameters example: $ python train.py --hyp {evolve_yaml}')
+        LOGGER.info(f'Hyperparameter evolution finished\n'
+                    f"Results saved to {colorstr('bold', save_dir)}\n"
+                    f'Use best hyperparameters example: $ python train.py --hyp {evolve_yaml}')
 
 
 def run(**kwargs):