Spaces:
Runtime error
Runtime error
File size: 1,506 Bytes
0b7b08a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import megengine as mge
import megengine.module as M
from models.yolo_fpn import YOLOFPN
from models.yolo_head import YOLOXHead
from models.yolo_pafpn import YOLOPAFPN
from models.yolox import YOLOX
def build_yolox(name="yolox-s"):
num_classes = 80
# value meaning: depth, width
param_dict = {
"yolox-nano": (0.33, 0.25),
"yolox-tiny": (0.33, 0.375),
"yolox-s": (0.33, 0.50),
"yolox-m": (0.67, 0.75),
"yolox-l": (1.0, 1.0),
"yolox-x": (1.33, 1.25),
}
if name == "yolov3":
depth = 1.0
width = 1.0
backbone = YOLOFPN()
head = YOLOXHead(num_classes, width, in_channels=[128, 256, 512], act="lrelu")
model = YOLOX(backbone, head)
else:
assert name in param_dict
kwargs = {}
depth, width = param_dict[name]
if name == "yolox-nano":
kwargs["depthwise"] = True
in_channels = [256, 512, 1024]
backbone = YOLOPAFPN(depth, width, in_channels=in_channels, **kwargs)
head = YOLOXHead(num_classes, width, in_channels=in_channels, **kwargs)
model = YOLOX(backbone, head)
for m in model.modules():
if isinstance(m, M.BatchNorm2d):
m.eps = 1e-3
return model
def build_and_load(weight_file, name="yolox-s"):
model = build_yolox(name)
model_weights = mge.load(weight_file)
model.load_state_dict(model_weights, strict=False)
return model
|