From e7269211ee130ff19b22d6e141f15e6dde261dde Mon Sep 17 00:00:00 2001 From: Kaiheng Weng <107049943+khwengXU@users.noreply.github.com> Date: Thu, 23 Jun 2022 16:11:04 +0800 Subject: [PATCH] feat: update minor (#27) --- deploy/ONNX/export_onnx.py | 2 +- yolov6/core/evaler.py | 13 ++++++------- yolov6/models/effidehead.py | 5 ++++- yolov6/models/yolo.py | 4 ++++ yolov6/utils/checkpoint.py | 1 + yolov6/utils/ema.py | 1 + yolov6/utils/envs.py | 1 + yolov6/utils/events.py | 3 +++ 8 files changed, 21 insertions(+), 9 deletions(-) diff --git a/deploy/ONNX/export_onnx.py b/deploy/ONNX/export_onnx.py index addd8ff6..6d5040ea 100644 --- a/deploy/ONNX/export_onnx.py +++ b/deploy/ONNX/export_onnx.py @@ -26,7 +26,7 @@ parser.add_argument('--batch-size', type=int, default=1, help='batch size') parser.add_argument('--half', action='store_true', help='FP16 half-precision export') parser.add_argument('--inplace', action='store_true', help='set Detect() inplace=True') - parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') + parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0, 1, 2, 3 or cpu') args = parser.parse_args() args.img_size *= 2 if len(args.img_size) == 1 else 1 # expand print(args) diff --git a/yolov6/core/evaler.py b/yolov6/core/evaler.py index 16b2e50e..ef617177 100644 --- a/yolov6/core/evaler.py +++ b/yolov6/core/evaler.py @@ -61,7 +61,6 @@ def init_model(self, model, weights, task): def init_data(self, dataloader, task): '''Initialize dataloader. - Returns a dataloader for task val or speed. ''' self.is_coco = isinstance(self.data.get('val'), str) and 'coco' in self.data['val'] # COCO dataset @@ -104,10 +103,10 @@ def predict_model(self, model, dataloader, task): return pred_results def eval_model(self, pred_results, model, dataloader, task): - '''Evaluate current model - For task speed, this function only evaluates the speed of model and output inference time. - For task val, this function evalutates the speed and also evaluates mAP by pycocotools, and then - returns inference time and mAP value. + '''Evaluate models + For task speed, this function only evaluates the speed of model and outputs inference time. + For task val, this function evalutates the speed and mAP by pycocotools, and returns + inference time and mAP value. ''' LOGGER.info(f'\nEvaluating speed.') self.eval_speed(task) @@ -145,7 +144,7 @@ def eval_model(self, pred_results, model, dataloader, task): return (0.0, 0.0) def eval_speed(self, task): - '''Evaluate the speed of model.''' + '''Evaluate model inference speed.''' if task != 'train': n_samples = self.speed_result[0].item() pre_time, inf_time, nms_time = 1000 * self.speed_result[1:].cpu().numpy() / n_samples @@ -215,7 +214,7 @@ def check_task(task): @staticmethod def reload_thres(conf_thres, iou_thres, task): - '''Sets conf and iou thres for task val/speed''' + '''Sets conf and iou threshold for task val/speed''' if task != 'train': if task == 'val': conf_thres = 0.001 diff --git a/yolov6/models/effidehead.py b/yolov6/models/effidehead.py index 4d9c1c75..4664f937 100644 --- a/yolov6/models/effidehead.py +++ b/yolov6/models/effidehead.py @@ -5,7 +5,10 @@ class Detect(nn.Module): - '''Efficient Decoupled Head''' + '''Efficient Decoupled Head + With hardware-aware degisn, the decoupled head is optimized with + hybridchannels methods. + ''' def __init__(self, num_classes=80, anchors=1, num_layers=3, inplace=True, head_layers=None): # detection layer super().__init__() assert head_layers is not None diff --git a/yolov6/models/yolo.py b/yolov6/models/yolo.py index b30b6b89..5d3d86be 100644 --- a/yolov6/models/yolo.py +++ b/yolov6/models/yolo.py @@ -10,6 +10,10 @@ class Model(nn.Module): + '''YOLOv6 model with backbone, neck and head. + The default parts are EfficientRep Backbone, Rep-PAN and + Efficient Decoupled Head. + ''' def __init__(self, config, channels=3, num_classes=None, anchors=None): # model, input channels, number of classes super().__init__() # Build network diff --git a/yolov6/utils/checkpoint.py b/yolov6/utils/checkpoint.py index 1cc6a333..3ce2ede4 100644 --- a/yolov6/utils/checkpoint.py +++ b/yolov6/utils/checkpoint.py @@ -7,6 +7,7 @@ from yolov6.utils.events import LOGGER from yolov6.utils.torch_utils import fuse_model + def load_state_dict(weights, model, map_location=None): """Load weights from checkpoint file, only assign weights those layers' name and shape are match.""" ckpt = torch.load(weights, map_location=map_location) diff --git a/yolov6/utils/ema.py b/yolov6/utils/ema.py index 3bdaf29d..104d97b3 100644 --- a/yolov6/utils/ema.py +++ b/yolov6/utils/ema.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn + class ModelEMA: """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models Keep a moving average of everything in the model state_dict (parameters and buffers). diff --git a/yolov6/utils/envs.py b/yolov6/utils/envs.py index 2e04d532..10159a94 100644 --- a/yolov6/utils/envs.py +++ b/yolov6/utils/envs.py @@ -8,6 +8,7 @@ import torch.backends.cudnn as cudnn from yolov6.utils.events import LOGGER + def get_envs(): """Get PyTorch needed environments from system envirionments.""" local_rank = int(os.getenv('LOCAL_RANK', -1)) diff --git a/yolov6/utils/events.py b/yolov6/utils/events.py index 3fee8c91..6a3dd509 100644 --- a/yolov6/utils/events.py +++ b/yolov6/utils/events.py @@ -5,14 +5,17 @@ import logging import shutil + def set_logging(name=None): rank = int(os.getenv('RANK', -1)) logging.basicConfig(format="%(message)s", level=logging.INFO if (rank in (-1, 0)) else logging.WARNING) return logging.getLogger(name) + LOGGER = set_logging(__name__) NCOLS = shutil.get_terminal_size().columns + def load_yaml(file_path): """Load data from yaml file.""" if isinstance(file_path, str):