Add initial model interface (#30)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ultralytics/yolo/utils/configs/defaults.yaml

@@ -1,53 +1,56 @@
 model: null
 data: null
-train:
-  epochs: 300
-  batch_size: 16
-  img_size: 640
-  nosave: False
-  cache: False # True/ram for ram, or disc
-  device: '' # cuda device, i.e. 0 or 0,1,2,3 or cpu
-  workers: 8
-  project: "ultralytics-yolo"
-  name: "exp" # TODO: make this informative, maybe exp{#number}_{datetime} ?
-  exist_ok: False
-  pretrained: False
-  optimizer: "Adam" # choices=['SGD', 'Adam', 'AdamW', 'RMSProp']
-  verbose: False
-  seed: 0
-  local_rank: -1
 
-hyps:
-  lr0: 0.001  # initial learning rate (SGD=1E-2, Adam=1E-3)
-  lrf: 0.01  # final OneCycleLR learning rate (lr0 * lrf)
-  momentum: 0.937  # SGD momentum/Adam beta1
-  weight_decay: 0.0005  # optimizer weight decay 5e-4
-  warmup_epochs: 3.0  # warmup epochs (fractions ok)
-  warmup_momentum: 0.8  # warmup initial momentum
-  warmup_bias_lr: 0.1  # warmup initial bias lr
-  box: 0.05  # box loss gain
-  cls: 0.5  # cls loss gain
-  cls_pw: 1.0  # cls BCELoss positive_weight
-  obj: 1.0  # obj loss gain (scale with pixels)
-  obj_pw: 1.0  # obj BCELoss positive_weight
-  iou_t: 0.20  # IoU training threshold
-  anchor_t: 4.0  # anchor-multiple threshold
-  # anchors: 3  # anchors per output layer (0 to ignore)
-  fl_gamma: 0.0  # focal loss gamma (efficientDet default gamma=1.5)
-  hsv_h: 0.015  # image HSV-Hue augmentation (fraction)
-  hsv_s: 0.7  # image HSV-Saturation augmentation (fraction)
-  hsv_v: 0.4  # image HSV-Value augmentation (fraction)
-  degrees: 0.0  # image rotation (+/- deg)
-  translate: 0.1  # image translation (+/- fraction)
-  scale: 0.5  # image scale (+/- gain)
-  shear: 0.0  # image shear (+/- deg)
-  perspective: 0.0  # image perspective (+/- fraction), range 0-0.001
-  flipud: 0.0  # image flip up-down (probability)
-  fliplr: 0.5  # image flip left-right (probability)
-  mosaic: 1.0  # image mosaic (probability)
-  mixup: 0.0  # image mixup (probability)
-  copy_paste: 0.0  # segment copy-paste (probability)
+# Training options
+epochs: 300
+batch_size: 16
+img_size: 640
+nosave: False
+cache: False # True/ram for ram, or disc
+device: '' # cuda device, i.e. 0 or 0,1,2,3 or cpu
+workers: 8
+project: "ultralytics-yolo"
+name: "exp" # TODO: make this informative, maybe exp{#number}_{datetime} ?
+exist_ok: False
+pretrained: False
+optimizer: "Adam" # choices=['SGD', 'Adam', 'AdamW', 'RMSProp']
+verbose: False
+seed: 0
+local_rank: -1
+#-----------------------------------#
 
+# Hyper-parameters
+lr0: 0.001  # initial learning rate (SGD=1E-2, Adam=1E-3)
+lrf: 0.01  # final OneCycleLR learning rate (lr0 * lrf)
+momentum: 0.937  # SGD momentum/Adam beta1
+weight_decay: 0.0005  # optimizer weight decay 5e-4
+warmup_epochs: 3.0  # warmup epochs (fractions ok)
+warmup_momentum: 0.8  # warmup initial momentum
+warmup_bias_lr: 0.1  # warmup initial bias lr
+box: 0.05  # box loss gain
+cls: 0.5  # cls loss gain
+cls_pw: 1.0  # cls BCELoss positive_weight
+obj: 1.0  # obj loss gain (scale with pixels)
+obj_pw: 1.0  # obj BCELoss positive_weight
+iou_t: 0.20  # IoU training threshold
+anchor_t: 4.0  # anchor-multiple threshold
+# anchors: 3  # anchors per output layer (0 to ignore)
+fl_gamma: 0.0  # focal loss gamma (efficientDet default gamma=1.5)
+hsv_h: 0.015  # image HSV-Hue augmentation (fraction)
+hsv_s: 0.7  # image HSV-Saturation augmentation (fraction)
+hsv_v: 0.4  # image HSV-Value augmentation (fraction)
+degrees: 0.0  # image rotation (+/- deg)
+translate: 0.1  # image translation (+/- fraction)
+scale: 0.5  # image scale (+/- gain)
+shear: 0.0  # image shear (+/- deg)
+perspective: 0.0  # image perspective (+/- fraction), range 0-0.001
+flipud: 0.0  # image flip up-down (probability)
+fliplr: 0.5  # image flip left-right (probability)
+mosaic: 1.0  # image mosaic (probability)
+mixup: 0.0  # image mixup (probability)
+copy_paste: 0.0  # segment copy-paste (probability)
+
+# Hydra configs -------------------------------------
 # to disable hydra directory creation
 hydra:
   output_subdir: null

ultralytics/yolo/utils/modeling/tasks.py

@@ -8,7 +8,8 @@ from ultralytics.yolo.utils import LOGGER
 from ultralytics.yolo.utils.anchors import check_anchor_order
 from ultralytics.yolo.utils.modeling import parse_model
 from ultralytics.yolo.utils.modeling.modules import *
-from ultralytics.yolo.utils.torch_utils import fuse_conv_and_bn, initialize_weights, model_info, scale_img, time_sync
+from ultralytics.yolo.utils.torch_utils import (fuse_conv_and_bn, initialize_weights, intersect_state_dicts, model_info,
+                                                scale_img, time_sync)
 
 
 class BaseModel(nn.Module):
@@ -67,6 +68,10 @@ class BaseModel(nn.Module):
                 m.anchor_grid = list(map(fn, m.anchor_grid))
         return self
 
+    def load(self, weights):
+        # Force all tasks implement this function
+        raise NotImplementedError("This function needs to be implemented by derived classes!")
+
 
 class DetectionModel(BaseModel):
     # YOLO detection model
@@ -166,6 +171,12 @@ class DetectionModel(BaseModel):
             b.data[:, 5:5 + m.nc] += math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum())  # cls
             mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
 
+    def load(self, weights):
+        ckpt = torch.load(weights, map_location='cpu')  # load checkpoint to CPU to avoid CUDA memory leak
+        csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
+        csd = intersect_state_dicts(csd, self.state_dict())  # intersect
+        self.load_state_dict(csd, strict=False)  # load
+
 
 class SegmentationModel(DetectionModel):
     # YOLOv5 segmentation model
@@ -197,3 +208,9 @@ class ClassificationModel(BaseModel):
     def _from_yaml(self, cfg):
         # Create a YOLOv5 classification model from a *.yaml file
         self.model = None
+
+    def load(self, weights):
+        ckpt = torch.load(weights, map_location='cpu')  # load checkpoint to CPU to avoid CUDA memory leak
+        csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
+        csd = intersect_state_dicts(csd, self.state_dict())  # intersect
+        self.load_state_dict(csd, strict=False)  # load

ultralytics/yolo/utils/torch_utils.py

@@ -174,3 +174,8 @@ def smart_inference_mode(torch_1_9=check_version(torch.__version__, '1.9.0')):
         return (torch.inference_mode if torch_1_9 else torch.no_grad)()(fn)
 
     return decorate
+
+
+def intersect_state_dicts(da, db, exclude=()):
+    # Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
+    return {k: v for k, v in da.items() if k in db and all(x not in k for x in exclude) and v.shape == db[k].shape}