【PyTorch】PyTorch预训练模型缓存位置迁移,也可拓展应用于其他文件的迁移
目录
前言:
一、具体实现:
二、关键技术解析
路径动态拼接
安全目录创建
环境变量魔法
迁移条件检查
三、代码实现:
前言:
当模型文件下载到本地c盘的默认路径时,可用以下代码的形式进行文件位置的迁移。
一、具体实现:
以下Python脚本完美解决上述问题,核心流程如图:
二、关键技术解析
-
路径动态拼接
target_path = os.path.join(project_dir, "xxxx.pth")
使用
os.path.join
确保跨平台路径兼容性 -
安全目录创建
os.makedirs(project_dir, exist_ok=True)
exist_ok=True
参数避免目录已存在时报错 -
环境变量魔法
os.environ['TORCH_HOME'] = os.path.dirname(project_dir)
通过设置
TORCH_HOME
永久改变缓存根目录(进程级生效) -
迁移条件检查
-
源文件存在且目标文件不存在 → 执行迁移
-
目标文件已存在 → 跳过迁移
-
源文件不存在 → 提示下载
-
三、代码实现:
import os
import shutil
import torch
import torchvision.models as models# 源文件路径(你的缓存目录)
source_path = r"C:\Users\14844\.cache\torch\hub\checkpoints\xxxxx.pth"# 目标路径(你的项目目录)
project_dir = r"E:\Pycharm\project\xxxx\models"
target_path = os.path.join(project_dir, "xxxxx.pth")# 创建目标目录(如果不存在)
os.makedirs(project_dir, exist_ok=True)# 移动模型文件
if os.path.exists(source_path) and not os.path.exists(target_path):print(f"正在将模型从 {source_path} 移动到 {target_path}...")shutil.move(source_path, target_path)print("模型移动完成!")
elif os.path.exists(target_path):print(f"模型已存在于 {target_path},无需移动。")
else:print(f"源模型文件不存在:{source_path}")print("请确保模型已正确下载。")# 更新环境变量,指向新的模型位置
os.environ['TORCH_HOME'] = os.path.dirname(project_dir)# 测试加载模型
try:print("\n正在测试加载模型...")vgg = models.vgg19(pretrained=True)print("模型加载成功!程序可以正常运行。")
except Exception as e:print(f"模型加载失败:{e}")print("请检查路径是否正确或尝试重新下载模型。")