Github项目仓库中官方项目提供的预训练模型下载很慢？没有办法下载下载？得一直盯着看进度下载失败了手动点击重新下载？基于python来开发实现自动化预训练模型权重下载存储模块实现预训练模型权重下载

因为项目开发的原因，在实际的项目中经常会需要基于真实项目中产生的数据来进行模型的开发训练，但是考虑到实际高质量标注数据集的高成本和稀缺性往往不会重头开始开发训练模型，一般的方法都是会基于官方提供的预训练权重模型来进行迁移训练，对于提升模型的效果和收敛速度，降低训练成本都有一定的帮助。

在今年的项目中又以目标检测系列的模型开发居多，从yolov3到yolov8，YOLO系列的版本迭代可以说是非常之快了，以yolov8为例，项目地址在这里，如下所示：

目前已有将近1.8w的star量。

yolov8项目本身同时支持：分类识别、目标检测、图像分割、姿态估计几种不同的任务类型的，对于不同的任务类型官方抖开放了对应的预训练权重可以直接使用，如下：

【分类识别】

Model	size (pixels)	acc top1	acc top5	Speed CPU ONNX (ms)	Speed A100 TensorRT (ms)	params (M)	FLOPs (B) at 640
YOLOv8n-cls	224	66.6	87.0	12.9	0.31	2.7	4.3
YOLOv8s-cls	224	72.3	91.1	23.4	0.35	6.4	13.5
YOLOv8m-cls	224	76.4	93.2	85.4	0.62	17.0	42.7
YOLOv8l-cls	224	78.0	94.1	163.0	0.87	37.5	99.7
YOLOv8x-cls	224	78.4	94.3	232.0	1.01	57.4	154.8

【目标检测】

Model	size (pixels)	mAPval 50-95	Speed CPU ONNX (ms)	Speed A100 TensorRT (ms)	params (M)	FLOPs (B)
YOLOv8n	640	37.3	80.4	0.99	3.2	8.7
YOLOv8s	640	44.9	128.4	1.20	11.2	28.6
YOLOv8m	640	50.2	234.7	1.83	25.9	78.9
YOLOv8l	640	52.9	375.2	2.39	43.7	165.2
YOLOv8x	640	53.9	479.1	3.53	68.2	257.8

目标检测还有另外一套是基于Open Image V7数据集训练出来的权重文件，如下所示：

Model	size (pixels)	mAPval 50-95	Speed CPU ONNX (ms)	Speed A100 TensorRT (ms)	params (M)	FLOPs (B)
YOLOv8n	640	18.4	142.4	1.21	3.5	10.5
YOLOv8s	640	27.7	183.1	1.40	11.4	29.7
YOLOv8m	640	33.6	408.5	2.26	26.2	80.6
YOLOv8l	640	34.9	596.9	2.43	44.1	167.4
YOLOv8x	640	36.3	860.6	3.56	68.7	260.6

【图像分割】

Model	size (pixels)	mAPbox 50-95	mAPmask 50-95	Speed CPU ONNX (ms)	Speed A100 TensorRT (ms)	params (M)	FLOPs (B)
YOLOv8n-seg	640	36.7	30.5	96.1	1.21	3.4	12.6
YOLOv8s-seg	640	44.6	36.8	155.7	1.47	11.8	42.6
YOLOv8m-seg	640	49.9	40.8	317.0	2.18	27.3	110.2
YOLOv8l-seg	640	52.3	42.6	572.4	2.79	46.0	220.5
YOLOv8x-seg	640	53.4	43.4	712.1	4.02	71.8	344.1

【姿态估计】

Model	size (pixels)	mAPpose 50-95	mAPpose 50	Speed CPU ONNX (ms)	Speed A100 TensorRT (ms)	params (M)	FLOPs (B)
YOLOv8n-pose	640	50.4	80.1	131.8	1.18	3.3	9.2
YOLOv8s-pose	640	60.0	86.2	233.2	1.42	11.6	30.2
YOLOv8m-pose	640	65.0	88.8	456.3	2.00	26.4	81.0
YOLOv8l-pose	640	67.6	90.0	784.5	2.59	44.4	168.6
YOLOv8x-pose	640	69.2	90.2	1607.1	3.73	69.4	263.2
YOLOv8x-pose-p6	1280	71.6	91.2	4088.7	10.04	99.1	1066.4

在我们想要使用预训练的时候，一般比较建议的方式是预先讲对应的预训练的模型下载到自己的本地项目中去然后将预训练模型的权重路径修改为本地路径这样会省略掉自动从网上下载的步骤，因为很多客观因素，联网下载失败率是很高的。

但是这里的问题是我们想要将预训练的模型预先下载到本地目录的话也是会经常失败的，甚至迟迟无法下载下来的，如下：

以目标检测任务yolov8x.pt权重下载为例：

这让人窒息的下载时长真的是无法接受，最无奈的是，这种下载状况大概率是等不到两小时就会下载失败了，还需要你得一直关注着文件是否下载成功了，如果没有下载成功就需要手动点击操作来重新下载，这是非常不方便的。

为了应对这种问题，这里我们开发实现了可以直接使用的自动预训练权重下载模块，可以帮助我们自动实现所需要预训练权重文件的下载，核心实现如下所示：

#!usr/bin/env python
# encoding:utf-8
from __future__ import division


"""
功能： 自动预训练模型权重下载模块
"""


import os
import time
import logging
import subprocess
import urllib
import requests
import torch
import random
from pathlib import Path



def autoDownload(file_path="all.txt", saveDir="weights/yolov8/"):
    """
    自动下载
    """
    if not os.path.exists(saveDir):
        os.makedirs(saveDir)
    while True:
        with open(file_path) as f:
            alls = [one.strip() for one in f.readlines() if one.strip()]
        count = 0
        for one_url in alls:
            print("downloading url: ", one_url)
            try:
                one_file = saveDir + one_url.split("/")[-1].strip()
                print("one_file: ", one_file)
                if not os.path.exists(one_file):
                    torch.hub.download_url_to_file(one_url, str(one_file))
                else:
                    print("当前预训练模型权重文件已经下载完成!")
                    count += 1
            except Exception as e:
                print("Exception: ", e)
                time.sleep(random.randint(1, 5))
        print("finished num: ", count)
        if count == len(alls):
            break



if __name__ == "__main__":


    autoDownload(file_path="all.txt", saveDir="weights/")

这里最大的好处就是只要你启动之后，他就会无休无止的帮你去完成你所需要的模型权重文件的下载任务，报错信息都可以忽略，如下：

这是我本地下载成功的截图，下载结果如下所示：

这里我们需要人工做的就是将自己需要下载的权重链接存放早alls.txt文件中，这里我是要下载yolov8所有的预训练权重文件，我的alls.txt如下所示：

https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-seg.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s-seg.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m-seg.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l-seg.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x-seg.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-pose.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s-pose.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m-pose.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l-pose.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x-pose.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x-pose-p6.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-cls.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s-cls.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m-cls.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l-cls.pt
https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x-cls.pt

需要的话可以直接复制拿去使用即可，这里我也是只是以yolov8为例讲解实现了全部预训练权重模型的下载存储实现，其他系列的模型也是一样的流程，感兴趣的话可以拿去试试！