diff --git a/models/tf.py b/models/tf.py index 62ba3ebf0782..17cca1e54fcf 100644 --- a/models/tf.py +++ b/models/tf.py @@ -380,9 +380,12 @@ def call(self, inputs): def parse_model(d, ch, model, imgsz): # model_dict, input_channels(3) LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}") - anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] + anchors, nc, gd, gw, ch_mul = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'], d.get( + 'channel_multiple') na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors no = na * (nc + 5) # number of outputs = anchors * (classes + 5) + if not ch_mul: + ch_mul = 8 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 @@ -399,7 +402,7 @@ def parse_model(d, ch, model, imgsz): # model_dict, input_channels(3) nn.Conv2d, Conv, DWConv, DWConvTranspose2d, Bottleneck, SPP, SPPF, MixConv2d, Focus, CrossConv, BottleneckCSP, C3, C3x]: c1, c2 = ch[f], args[0] - c2 = make_divisible(c2 * gw, 8) if c2 != no else c2 + c2 = make_divisible(c2 * gw, ch_mul) if c2 != no else c2 args = [c1, c2, *args[1:]] if m in [BottleneckCSP, C3, C3x]: @@ -414,7 +417,7 @@ def parse_model(d, ch, model, imgsz): # model_dict, input_channels(3) if isinstance(args[1], int): # number of anchors args[1] = [list(range(args[1] * 2))] * len(f) if m is Segment: - args[3] = make_divisible(args[3] * gw, 8) + args[3] = make_divisible(args[3] * gw, ch_mul) args.append(imgsz) else: c2 = ch[f] diff --git a/models/yolo.py b/models/yolo.py index 4f4d567bec73..9cd3cddc25c3 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -299,10 +299,13 @@ def _from_yaml(self, cfg): def parse_model(d, ch): # model_dict, input_channels(3) # Parse a YOLOv5 model.yaml dictionary LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}") - anchors, nc, gd, gw, act = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'], d.get('activation') + anchors, nc, gd, gw, act, ch_mul = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'], d.get( + 'activation'), d.get('channel_multiple') if act: Conv.default_act = eval(act) # redefine default activation, i.e. Conv.default_act = nn.SiLU() LOGGER.info(f"{colorstr('activation:')} {act}") # print + if not ch_mul: + ch_mul = 8 na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors no = na * (nc + 5) # number of outputs = anchors * (classes + 5) @@ -319,7 +322,7 @@ def parse_model(d, ch): # model_dict, input_channels(3) BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x}: c1, c2 = ch[f], args[0] if c2 != no: # if not output - c2 = make_divisible(c2 * gw, 8) + c2 = make_divisible(c2 * gw, ch_mul) args = [c1, c2, *args[1:]] if m in {BottleneckCSP, C3, C3TR, C3Ghost, C3x}: @@ -335,7 +338,7 @@ def parse_model(d, ch): # model_dict, input_channels(3) if isinstance(args[1], int): # number of anchors args[1] = [list(range(args[1] * 2))] * len(f) if m is Segment: - args[3] = make_divisible(args[3] * gw, 8) + args[3] = make_divisible(args[3] * gw, ch_mul) elif m is Contract: c2 = ch[f] * args[0] ** 2 elif m is Expand: