feat: update minor (#27)

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)

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

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

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

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)

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

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

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