ultralytics 8.0.55
unified YOLOv8 model YAMLs (#1475)
This commit is contained in:
@ -170,7 +170,7 @@ class DetectionModel(BaseModel):
|
||||
# YOLOv8 detection model
|
||||
def __init__(self, cfg='yolov8n.yaml', ch=3, nc=None, verbose=True): # model, input channels, number of classes
|
||||
super().__init__()
|
||||
self.yaml = cfg if isinstance(cfg, dict) else yaml_load(check_yaml(cfg), append_filename=True) # cfg dict
|
||||
self.yaml = cfg if isinstance(cfg, dict) else yaml_model_load(cfg) # cfg dict
|
||||
|
||||
# Define model
|
||||
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
|
||||
@ -277,7 +277,8 @@ class ClassificationModel(BaseModel):
|
||||
self.nc = nc
|
||||
|
||||
def _from_yaml(self, cfg, ch, nc, verbose):
|
||||
self.yaml = cfg if isinstance(cfg, dict) else yaml_load(check_yaml(cfg), append_filename=True) # cfg dict
|
||||
self.yaml = cfg if isinstance(cfg, dict) else yaml_model_load(cfg) # cfg dict
|
||||
|
||||
# Define model
|
||||
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
|
||||
if nc and nc != self.yaml['nc']:
|
||||
@ -418,30 +419,42 @@ def attempt_load_one_weight(weight, device=None, inplace=True, fuse=False):
|
||||
|
||||
|
||||
def parse_model(d, ch, verbose=True): # model_dict, input_channels(3)
|
||||
# Parse a YOLO model.yaml dictionary
|
||||
if verbose:
|
||||
LOGGER.info(f"\n{'':>3}{'from':>20}{'n':>3}{'params':>10} {'module':<45}{'arguments':<30}")
|
||||
nc, gd, gw, act = d['nc'], d['depth_multiple'], d['width_multiple'], d.get('activation')
|
||||
# Parse a YOLO model.yaml dictionary into a PyTorch model
|
||||
import ast
|
||||
|
||||
# Args
|
||||
max_channels = float('inf')
|
||||
nc, act, scales = (d.get(x) for x in ('nc', 'act', 'scales'))
|
||||
depth, width = (d.get(x, 1.0) for x in ('depth_multiple', 'width_multiple'))
|
||||
if scales:
|
||||
scale = d.get('scale')
|
||||
if not scale:
|
||||
scale = tuple(scales.keys())[0]
|
||||
LOGGER.warning(f"WARNING ⚠️ no model scale passed. Assuming scale='{scale}'.")
|
||||
depth, width, max_channels = scales[scale]
|
||||
|
||||
if act:
|
||||
Conv.default_act = eval(act) # redefine default activation, i.e. Conv.default_act = nn.SiLU()
|
||||
if verbose:
|
||||
LOGGER.info(f"{colorstr('activation:')} {act}") # print
|
||||
|
||||
if verbose:
|
||||
LOGGER.info(f"\n{'':>3}{'from':>20}{'n':>3}{'params':>10} {'module':<45}{'arguments':<30}")
|
||||
ch = [ch]
|
||||
layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out
|
||||
for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args
|
||||
m = getattr(torch.nn, m[3:]) if 'nn.' in m else globals()[m] # get module
|
||||
for j, a in enumerate(args):
|
||||
# TODO: re-implement with eval() removal if possible
|
||||
# args[j] = (locals()[a] if a in locals() else ast.literal_eval(a)) if isinstance(a, str) else a
|
||||
with contextlib.suppress(NameError):
|
||||
args[j] = eval(a) if isinstance(a, str) else a # eval strings
|
||||
if isinstance(a, str):
|
||||
with contextlib.suppress(ValueError):
|
||||
args[j] = locals()[a] if a in locals() else ast.literal_eval(a)
|
||||
|
||||
n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain
|
||||
n = n_ = max(round(n * depth), 1) if n > 1 else n # depth gain
|
||||
if m in (Classify, Conv, ConvTranspose, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, Focus,
|
||||
BottleneckCSP, C1, C2, C2f, C3, C3TR, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x):
|
||||
c1, c2 = ch[f], args[0]
|
||||
if c2 != nc: # if c2 not equal to number of classes (i.e. for Classify() output)
|
||||
c2 = make_divisible(c2 * gw, 8)
|
||||
c2 = make_divisible(min(c2, max_channels) * width, 8)
|
||||
|
||||
args = [c1, c2, *args[1:]]
|
||||
if m in (BottleneckCSP, C1, C2, C2f, C3, C3TR, C3Ghost, C3x):
|
||||
@ -454,7 +467,7 @@ def parse_model(d, ch, verbose=True): # model_dict, input_channels(3)
|
||||
elif m in (Detect, Segment):
|
||||
args.append([ch[x] for x in f])
|
||||
if m is Segment:
|
||||
args[2] = make_divisible(args[2] * gw, 8)
|
||||
args[2] = make_divisible(min(args[2], max_channels) * width, 8)
|
||||
else:
|
||||
c2 = ch[f]
|
||||
|
||||
@ -472,6 +485,41 @@ def parse_model(d, ch, verbose=True): # model_dict, input_channels(3)
|
||||
return nn.Sequential(*layers), sorted(save)
|
||||
|
||||
|
||||
def yaml_model_load(path):
|
||||
import re
|
||||
|
||||
path = Path(path)
|
||||
if path.stem in (f'yolov{d}{x}6' for x in 'nsmlx' for d in (5, 8)):
|
||||
new_stem = re.sub(r'(\d+)([nslmx])6(.+)?$', r'\1\2-p6\3', path.stem)
|
||||
LOGGER.warning(f'WARNING ⚠️ Ultralytics YOLO P6 models now use -p6 suffix. Renaming {path.stem} to {new_stem}.')
|
||||
path = path.with_stem(new_stem)
|
||||
|
||||
unified_path = re.sub(r'(\d+)([nslmx])(.+)?$', r'\1\3', str(path)) # i.e. yolov8x.yaml -> yolov8.yaml
|
||||
yaml_file = check_yaml(unified_path, hard=False) or check_yaml(path)
|
||||
d = yaml_load(yaml_file) # model dict
|
||||
d['scale'] = guess_model_scale(path)
|
||||
d['yaml_file'] = str(path)
|
||||
return d
|
||||
|
||||
|
||||
def guess_model_scale(model_path):
|
||||
"""
|
||||
Takes a path to a YOLO model's YAML file as input and extracts the size character of the model's scale.
|
||||
The function uses regular expression matching to find the pattern of the model scale in the YAML file name,
|
||||
which is denoted by n, s, m, l, or x. The function returns the size character of the model scale as a string.
|
||||
|
||||
Args:
|
||||
model_path (str or Path): The path to the YOLO model's YAML file.
|
||||
|
||||
Returns:
|
||||
(str): The size character of the model's scale, which can be n, s, m, l, or x.
|
||||
"""
|
||||
with contextlib.suppress(AttributeError):
|
||||
import re
|
||||
return re.search(r'yolov\d+([nslmx])', Path(model_path).stem).group(1) # n, s, m, l, or x
|
||||
return ''
|
||||
|
||||
|
||||
def guess_model_task(model):
|
||||
"""
|
||||
Guess the task of a PyTorch model from its architecture or configuration.
|
||||
|
Reference in New Issue
Block a user