From 442c79a0da6b9d51b90398e98f554f7a87b97696 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 28 Mar 2022 12:45:28 +0200 Subject: [PATCH 1/2] yolo.py profiling updates --- models/yolo.py | 26 ++++++++++++-------------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/models/yolo.py b/models/yolo.py index e88db79ca8c7..d9480a47e9ab 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -25,7 +25,8 @@ from utils.autoanchor import check_anchor_order from utils.general import LOGGER, check_version, check_yaml, make_divisible, print_args from utils.plots import feature_visualization -from utils.torch_utils import fuse_conv_and_bn, initialize_weights, model_info, scale_img, select_device, time_sync +from utils.torch_utils import fuse_conv_and_bn, initialize_weights, model_info, profile, scale_img, select_device, \ + time_sync try: import thop # for FLOPs computation @@ -300,8 +301,10 @@ def parse_model(d, ch): # model_dict, input_channels(3) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') + parser.add_argument('--batch-size', type=int, default=1, help='total batch size for all GPUs') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--profile', action='store_true', help='profile model speed') + parser.add_argument('--line-profile', action='store_true', help='profile model speed layer by layer') parser.add_argument('--test', action='store_true', help='test all yolo*.yaml') opt = parser.parse_args() opt.cfg = check_yaml(opt.cfg) # check YAML @@ -309,24 +312,19 @@ def parse_model(d, ch): # model_dict, input_channels(3) device = select_device(opt.device) # Create model + im = torch.rand(opt.batch_size, 3, 640, 640).to(device) model = Model(opt.cfg).to(device) - # Profile - if opt.profile: - model.eval().fuse() - img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) - y = model(img, profile=True) + # Options + if opt.line_profile: # profile layer by layer + _ = model(im, profile=True) - # Test all models - if opt.test: + elif opt.profile: # profile forward-backward + results = profile(input=im, ops=[model], n=3) + + elif opt.test: # test all models for cfg in Path(ROOT / 'models').rglob('yolo*.yaml'): try: _ = Model(cfg) except Exception as e: print(f'Error in {cfg}: {e}') - - # Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898) - # from torch.utils.tensorboard import SummaryWriter - # tb_writer = SummaryWriter('.') - # 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 From 336ac80b6c8015855664ad9513606d1fa5ac9fc5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 28 Mar 2022 10:45:42 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- models/yolo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/yolo.py b/models/yolo.py index d9480a47e9ab..81ab539deffa 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -25,8 +25,8 @@ from utils.autoanchor import check_anchor_order from utils.general import LOGGER, check_version, check_yaml, make_divisible, print_args from utils.plots import feature_visualization -from utils.torch_utils import fuse_conv_and_bn, initialize_weights, model_info, profile, scale_img, select_device, \ - time_sync +from utils.torch_utils import (fuse_conv_and_bn, initialize_weights, model_info, profile, scale_img, select_device, + time_sync) try: import thop # for FLOPs computation