【保姆级教程】【YOLOv8替换主干网络】【1】使用efficientViT替换YOLOV8主干网络结构(2)https://developer.aliyun.com/article/1536651
第2步–修改tasks.py中的相关内容
parse_model函数修改
修改ultralytics/nn/tasks.py
中的parse_model
函数,修改后完整代码如下:
def parse_model(d, ch, verbose=True): # model_dict, input_channels(3) """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', 'activation', 'scales')) depth, width, kpt_shape = (d.get(x, 1.0) for x in ('depth_multiple', 'width_multiple', 'kpt_shape')) 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 is_backbone = False for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args try: if m == 'node_mode': m = d[m] if len(args) > 0: if args[0] == 'head_channel': args[0] = int(d[args[0]]) t = m m = getattr(torch.nn, m[3:]) if 'nn.' in m else globals()[m] # get module except: pass for j, a in enumerate(args): if isinstance(a, str): with contextlib.suppress(ValueError): try: args[j] = locals()[a] if a in locals() else ast.literal_eval(a) except: args[j] = a 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, RepC3): 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(min(c2, max_channels) * width, 8) args = [c1, c2, *args[1:]] if m in (BottleneckCSP, C1, C2, C2f, C3, C3TR, C3Ghost, C3x, RepC3): args.insert(2, n) # number of repeats n = 1 elif m is AIFI: args = [ch[f], *args] elif m in (HGStem, HGBlock): c1, cm, c2 = ch[f], args[0], args[1] args = [c1, cm, c2, *args[2:]] if m is HGBlock: args.insert(4, n) # number of repeats n = 1 elif m is ResNetLayer: c2 = args[1] if args[3] else args[1] * 4 elif m is nn.BatchNorm2d: args = [ch[f]] elif m is Concat: c2 = sum(ch[x] for x in f) elif m in (Detect, Segment, Pose): args.append([ch[x] for x in f]) if m is Segment: args[2] = make_divisible(min(args[2], max_channels) * width, 8) elif m is RTDETRDecoder: # special case, channels arg must be passed in index 1 args.insert(1, [ch[x] for x in f]) elif m in {MHSA, ShuffleAttention}: args = [ch[f], *args] elif m in {EfficientViT_M0, EfficientViT_M1, EfficientViT_M2, EfficientViT_M3, EfficientViT_M4, EfficientViT_M5}: m = m(*args) c2 = m.channel else: c2 = ch[f] if isinstance(c2, list): is_backbone = True m_ = m m_.backbone = True else: m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module t = str(m)[8:-2].replace('__main__.', '') # module type m.np = sum(x.numel() for x in m_.parameters()) # number params m_.i, m_.f, m_.type = i + 4 if is_backbone else i, f, t # attach index, 'from' index, type if verbose: LOGGER.info(f'{i:>3}{str(f):>20}{n_:>3}{m.np:10.0f} {t:<45}{str(args):<30}') # print save.extend(x % (i + 4 if is_backbone else i) for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist layers.append(m_) if i == 0: ch = [] if isinstance(c2, list): ch.extend(c2) for _ in range(5 - len(ch)): ch.insert(0, 0) else: ch.append(c2) return nn.Sequential(*layers), sorted(save)
parse_model修改的详细内容对比
- 将
efficientVit.py
中的all
参数中的函数名,写入tasks.py
的判断分支中。
新建if判断分支,添加如下内容:
elif m in {efficientvit_b0, efficientvit_b1, efficientvit_b2, efficientvit_b3}: m = m(*args) c2 = m.channel
2.修改下图解析部分代码1,如下图:
修改前:
修改后:
代码如下:
is_backbone = False for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args try: if m == 'node_mode': m = d[m] if len(args) > 0: if args[0] == 'head_channel': args[0] = int(d[args[0]]) t = m m = getattr(torch.nn, m[3:]) if 'nn.' in m else globals()[m] # get module except: pass for j, a in enumerate(args): if isinstance(a, str): with contextlib.suppress(ValueError): try: args[j] = locals()[a] if a in locals() else ast.literal_eval(a) except: args[j] = a
3.修改下面截图中的部分代码2
修改前:
修改后:
代码如下:
if isinstance(c2, list): is_backbone = True m_ = m m_.backbone = True else: m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module t = str(m)[8:-2].replace('__main__.', '') # module type m.np = sum(x.numel() for x in m_.parameters()) # number params m_.i, m_.f, m_.type = i + 4 if is_backbone else i, f, t # attach index, 'from' index, type if verbose: LOGGER.info(f'{i:>3}{str(f):>20}{n_:>3}{m.np:10.0f} {t:<45}{str(args):<30}') # print save.extend(x % (i + 4 if is_backbone else i) for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist layers.append(m_)
4.修改下面截图部分代码
修改前:
修改后:
修改代码为:
if isinstance(c2, list): ch.extend(c2) for _ in range(5 - len(ch)): ch.insert(0, 0) else: ch.append(c2)
【保姆级教程】【YOLOv8替换主干网络】【1】使用efficientViT替换YOLOV8主干网络结构(4)https://developer.aliyun.com/article/1536655