diff --git a/mindyolo/models/heads/yolov5_head.py b/mindyolo/models/heads/yolov5_head.py index b4e2b89f..0873d600 100644 --- a/mindyolo/models/heads/yolov5_head.py +++ b/mindyolo/models/heads/yolov5_head.py @@ -1,3 +1,4 @@ +import math import numpy as np import mindspore as ms @@ -48,6 +49,16 @@ def construct(self, x): # return outs return outs if self.training else (ops.concat(z, 1), outs) + def initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency + # https://arxiv.org/abs/1708.02002 section 3.3 + m = self + for mi, s in zip(m.m, m.stride): # from + s = s.asnumpy() + b = mi.bias.view(m.na, -1).asnumpy() # conv.bias(255) to (3,85) + b[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) + b[:, 5:] += math.log(0.6 / (m.nc - 0.999999)) if cf is None else np.log(cf / cf.sum()) # cls + mi.bias = ops.assign(mi.bias, Tensor(b, ms.float32).view(-1)) + @staticmethod def _make_grid(nx=20, ny=20, dtype=ms.float32): # FIXME: Not supported on a specific model of machine diff --git a/mindyolo/models/yolov3.py b/mindyolo/models/yolov3.py index 247064eb..9f87d2fe 100644 --- a/mindyolo/models/yolov3.py +++ b/mindyolo/models/yolov3.py @@ -3,9 +3,9 @@ import mindspore as ms from mindspore import Tensor, nn -from .heads.yolov3_head import YOLOv3Head -from .model_factory import build_model_from_cfg -from .registry import register_model +from mindyolo.models.heads.yolov3_head import YOLOv3Head +from mindyolo.models.model_factory import build_model_from_cfg +from mindyolo.models.registry import register_model __all__ = ["YOLOv3", "yolov3"] @@ -50,16 +50,17 @@ def yolov3(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv3: if __name__ == "__main__": from mindyolo.models.model_factory import create_model - from mindyolo.utils.config import parse_config - - opt = parse_config() - model = create_model( - model_name="yolov3", - model_cfg=opt.net, - num_classes=opt.data.nc, - sync_bn=opt.sync_bn if hasattr(opt, "sync_bn") else False, + from mindyolo.utils.config import load_config, Config + + cfg, _, _ = load_config('../../configs/yolov3/yolov3.yaml') + cfg = Config(cfg) + network = create_model( + model_name=cfg.network.model_name, + model_cfg=cfg.network, + num_classes=cfg.data.nc, + sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False, ) x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32) - out = model(x) + out = network(x) out = out[0] if isinstance(out, (list, tuple)) else out print(f"Output shape is {[o.shape for o in out]}") diff --git a/mindyolo/models/yolov4.py b/mindyolo/models/yolov4.py index b70a261b..08748c5f 100644 --- a/mindyolo/models/yolov4.py +++ b/mindyolo/models/yolov4.py @@ -1,7 +1,10 @@ -import mindspore.nn as nn +import numpy as np -from .model_factory import build_model_from_cfg -from .registry import register_model +import mindspore as ms +from mindspore import Tensor, nn + +from mindyolo.models.model_factory import build_model_from_cfg +from mindyolo.models.registry import register_model __all__ = ["YOLOv4", "yolov4"] @@ -32,3 +35,21 @@ def yolov4(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv4: """Get yolov4 model.""" model = YOLOv4(cfg=cfg, in_channels=in_channels, num_classes=num_classes, **kwargs) return model + + +if __name__ == "__main__": + from mindyolo.models.model_factory import create_model + from mindyolo.utils.config import load_config, Config + + cfg, _, _ = load_config('../../configs/yolov4/yolov4.yaml') + cfg = Config(cfg) + network = create_model( + model_name=cfg.network.model_name, + model_cfg=cfg.network, + num_classes=cfg.data.nc, + sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False, + ) + x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32) + out = network(x) + out = out[0] if isinstance(out, (list, tuple)) else out + print(f"Output shape is {[o.shape for o in out]}") diff --git a/mindyolo/models/yolov5.py b/mindyolo/models/yolov5.py index 6fee9f00..cb8931e8 100644 --- a/mindyolo/models/yolov5.py +++ b/mindyolo/models/yolov5.py @@ -3,9 +3,9 @@ import mindspore as ms from mindspore import Tensor, nn -from .heads.yolov7_head import YOLOv7AuxHead, YOLOv7Head -from .model_factory import build_model_from_cfg -from .registry import register_model +from mindyolo.models.heads.yolov5_head import YOLOv5Head +from mindyolo.models.model_factory import build_model_from_cfg +from mindyolo.models.registry import register_model __all__ = ["YOLOv5", "yolov5"] @@ -37,10 +37,8 @@ def construct(self, x): def initialize_weights(self): # reset parameter for Detect Head m = self.model.model[-1] - if isinstance(m, YOLOv7Head): + if isinstance(m, YOLOv5Head): m.initialize_biases() - if isinstance(m, YOLOv7AuxHead): - m.initialize_aux_biases() @register_model @@ -55,16 +53,17 @@ def yolov5(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv5: if __name__ == "__main__": from mindyolo.models.model_factory import create_model - from mindyolo.utils.config import parse_config - - opt = parse_config() - model = create_model( - model_name="yolov5", - model_cfg=opt.net, - num_classes=opt.data.nc, - sync_bn=opt.sync_bn if hasattr(opt, "sync_bn") else False, + from mindyolo.utils.config import load_config, Config + + cfg, _, _ = load_config('../../configs/yolov5/yolov5s.yaml') + cfg = Config(cfg) + network = create_model( + model_name=cfg.network.model_name, + model_cfg=cfg.network, + num_classes=cfg.data.nc, + sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False, ) x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32) - out = model(x) + out = network(x) out = out[0] if isinstance(out, (list, tuple)) else out print(f"Output shape is {[o.shape for o in out]}") diff --git a/mindyolo/models/yolov7.py b/mindyolo/models/yolov7.py index 393b0b0e..a0b100aa 100644 --- a/mindyolo/models/yolov7.py +++ b/mindyolo/models/yolov7.py @@ -3,8 +3,8 @@ import mindspore as ms from mindspore import Tensor, nn -from .heads.yolov7_head import YOLOv7AuxHead, YOLOv7Head -from .model_factory import build_model_from_cfg +from mindyolo.models.heads.yolov7_head import YOLOv7AuxHead, YOLOv7Head +from mindyolo.models.model_factory import build_model_from_cfg from .registry import register_model __all__ = ["YOLOv7", "yolov7"] @@ -55,16 +55,17 @@ def yolov7(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv7: if __name__ == "__main__": from mindyolo.models.model_factory import create_model - from mindyolo.utils.config import parse_config - - opt = parse_config() - model = create_model( - model_name="yolov7", - model_cfg=opt.net, - num_classes=opt.data.nc, - sync_bn=opt.sync_bn if hasattr(opt, "sync_bn") else False, + from mindyolo.utils.config import load_config, Config + + cfg, _, _ = load_config('../../configs/yolov7/yolov7-tiny.yaml') + cfg = Config(cfg) + network = create_model( + model_name=cfg.network.model_name, + model_cfg=cfg.network, + num_classes=cfg.data.nc, + sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False, ) x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32) - out = model(x) + out = network(x) out = out[0] if isinstance(out, (list, tuple)) else out print(f"Output shape is {[o.shape for o in out]}") diff --git a/mindyolo/models/yolov8.py b/mindyolo/models/yolov8.py index befd1a70..f67301f3 100644 --- a/mindyolo/models/yolov8.py +++ b/mindyolo/models/yolov8.py @@ -3,9 +3,9 @@ import mindspore as ms from mindspore import Tensor, nn -from .heads.yolov8_head import YOLOv8Head -from .model_factory import build_model_from_cfg -from .registry import register_model +from mindyolo.models.heads.yolov8_head import YOLOv8Head +from mindyolo.models.model_factory import build_model_from_cfg +from mindyolo.models.registry import register_model __all__ = ["YOLOv8", "yolov8"] @@ -54,16 +54,17 @@ def yolov8(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOv8: if __name__ == "__main__": from mindyolo.models.model_factory import create_model - from mindyolo.utils.config import parse_config - - opt = parse_config() - model = create_model( - model_name="yolov8", - model_cfg=opt.net, - num_classes=opt.data.nc, - sync_bn=opt.sync_bn if hasattr(opt, "sync_bn") else False, + from mindyolo.utils.config import load_config, Config + + cfg, _, _ = load_config('../../configs/yolov8/yolov8s.yaml') + cfg = Config(cfg) + network = create_model( + model_name=cfg.network.model_name, + model_cfg=cfg.network, + num_classes=cfg.data.nc, + sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False, ) x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32) - out = model(x) + out = network(x) out = out[0] if isinstance(out, (list, tuple)) else out print(f"Output shape is {[o.shape for o in out]}") diff --git a/mindyolo/models/yolox.py b/mindyolo/models/yolox.py index f7463c1d..7dde8705 100644 --- a/mindyolo/models/yolox.py +++ b/mindyolo/models/yolox.py @@ -4,8 +4,8 @@ from mindspore import Tensor, nn from mindyolo.models.registry import register_model -from .heads import YOLOXHead -from .model_factory import build_model_from_cfg +from mindyolo.models.heads import YOLOXHead +from mindyolo.models.model_factory import build_model_from_cfg __all__ = ["YOLOX", "yolox"] @@ -45,3 +45,21 @@ def yolox(cfg, in_channels=3, num_classes=None, **kwargs) -> YOLOX: """Get yolox model.""" model = YOLOX(cfg, in_channels=in_channels, num_classes=num_classes, **kwargs) return model + + +if __name__ == "__main__": + from mindyolo.models.model_factory import create_model + from mindyolo.utils.config import load_config, Config + + cfg, _, _ = load_config('../../configs/yolox/yolox-s.yaml') + cfg = Config(cfg) + network = create_model( + model_name=cfg.network.model_name, + model_cfg=cfg.network, + num_classes=cfg.data.nc, + sync_bn=cfg.sync_bn if hasattr(cfg, "sync_bn") else False, + ) + x = Tensor(np.random.randn(1, 3, 640, 640), ms.float32) + out = network(x) + out = out[0] if isinstance(out, (list, tuple)) else out + print(f"Output shape is {[o.shape for o in out]}")