【学习bubbliiiing代码-1】下载github上神经网络权重

最新推荐文章于 2024-07-06 12:03:33 发布

寻找永不遗憾

最新推荐文章于 2024-07-06 12:03:33 发布

阅读量166

点赞数

分类专栏：深度学习基础知识文章标签：代码下载权重 python

本文链接：https://blog.csdn.net/weixin_45377629/article/details/131690848

版权

深度学习基础知识专栏收录该内容

67 篇文章 8 订阅

订阅专栏

本系列主要用于自我学习，参考的为bubbliiiing的代码

写一个优雅的下载权重的函数，如下：

def download_weights(phi, model_dir="./model_data"):
    import os
    from torch.hub import load_state_dict_from_url
    
    download_urls = {
        "n" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_n_backbone_weights.pth',
        "s" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_s_backbone_weights.pth',
        "m" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_m_backbone_weights.pth',
        "l" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_l_backbone_weights.pth',
        "x" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_x_backbone_weights.pth',
    }
    url = download_urls[phi]
    
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)
    load_state_dict_from_url(url, model_dir)


if __name__ == "__main__":
	phi = 's'
	download_weights(phi)