Yolov5-7.0图像分类算法修改Resnet18/50主干网络流程

网上大多数都是基于yolov5算法的目标检测网络进行修改主干网络,我最近在尝试图像分类算法,流程如下:

以resnet50为例

1、打开models下的common.py文件,添加下面的代码:

'''
模型:resnet50
'''


class resnet501(nn.Module):
    def __init__(self, ignore) -> None:
        super().__init__()
        model = models.resnet50(pretrained=True)
        modules = list(model.children())
        modules = modules[:6]
        self.model = nn.Sequential(*modules)

    def forward(self, x):
        return self.model(x)


class resnet502(nn.Module):
    def __init__(self, ignore) -> None:
        super().__init__()
        model = models.resnet50(pretrained=True)
        modules = list(model.children())
        modules = modules[6]
        self.model = nn.Sequential(*modules)

    def forward(self, x):
        return self.model(x)


class resnet503(nn.Module):
    def __init__(self, ignore) -> None:
        super().__init__()
        model = models.resnet50(pretrained=True)
        modules = list(model.children())
        modules = modules[7]
        self.model = nn.Sequential(*modules)

    def forward(self, x):
        return self. Model(x)

2、打开yolo.py文件,添加

        # 这行代码下添加
        elif m is Expand:
            c2 = ch[f] // args[0] ** 2
        #新添加
        elif m is resnet501 or m is resnet502 or m is resnet503:
            c2 = args[0]

3、新建resnet50.yaml文件,放在models文件夹下

# YOLOv5  by Ultralytics, GPL-3.0 license

# Parameters
nc: 2  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.25  # layer channel multiple
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5 v6.0 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, resnet501, [512]],  # 0
   [-1, 1, resnet502, [1024]],  # 1
   [-1, 1, resnet503, [2048]],  # 2
   [-1, 1, SPPF, [1024, 5]],  # 3
  ]

# YOLOv5 v6.0 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 1], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, C3, [512, False]],  # 7

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 0], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, C3, [256, False]],  # 11 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 7], 1, Concat, [1]],  # cat head P4
   [-1, 3, C3, [512, False]],  # 14 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 3], 1, Concat, [1]],  # cat head P5
   [-1, 3, C3, [1024, False]],  # 17 (P5/32-large)

   [[11, 14, 17], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

4、测试yolo.py文件

在if __name__ == "__main__":下修改这行代码,添加resnet50.yaml文件

    parser.add_argument('--cfg', type=str, default=ROOT/'resnet50.yaml', help='model.yaml')

运行:结果如下就是成功了

Yolov5-7.0图像分类算法修改Resnet18/50主干网络流程_第1张图片

如果不是,检查以上步骤,哪里出错了。。。。

5、修改classify文件夹下的train.py文件,首先修改在parse_opt(known = False)函数下,添加加修改如下:

    parser.add_argument('--model', type=str, default='resnet50', help='根路径下的预训练模型')
    parser.add_argument('--cfg', type=str, default=ROOT/'models/resnet50.yaml', help='根路径下的预训练模型')

然后,crtl+F定位 torchvision.models.__dict__ 这句代码,修改如下:

初始是这样的:
model = torchvision.models.__dict__[opt.model](weights='IMAGENET1K_V1' if pretrained else None)


修改之后:
model = torchvision.models.__dict__[opt.model](pretrained=opt.pretrained)

再次定位  ClassificationModel ,修改如下

初始是这样的:
model = ClassificationModel(model=model, nc=nc, cutoff=opt.cutoff or 10)  # convert to classification model


修改之后:
model = ClassificationModel(cfg = opt.cfg, model=model, nc=nc, cutoff=opt.cutoff or 10)  # convert to classification model

6、运行测试train.py文件,如下结果就是成功了

Yolov5-7.0图像分类算法修改Resnet18/50主干网络流程_第2张图片

你可能感兴趣的:(yolo,YOLO,图像分类,yolov5-7.0,深度学习,python)