模型权重也可能“有毒”?别再无脑 torch.load 了

部署运行你感兴趣的模型镜像

最近在训练 YOLO/Ultralytics 的同学可能都遇到过这样一行黄字提示:

C:\7zProject\SkinDiseaseTrain_v8\ultralytics\utils\torch_utils.py:521: FutureWarning: You are using torch.load with weights_only=False (the current default value), which uses the default pickle module implicitly.
It is possible to construct malicious pickle data which will execute arbitrary code during unpickling…

很多人第一反应是:“是不是模型坏了?”其实不是 bug,而是 PyTorch 在提醒你:模型权重文件如果来自不可信来源,真的可能“有毒”。

1.警告出现的原因

在 Ultralytics 内部,权重加载是这样写的:

x = torch.load(f, map_location=torch.device("cpu"))

因为没有传 weights_only=True,所以触发了 PyTorch 的 FutureWarning。

根本原因在于:
(1)weights_only=False(默认值)会调用 Python 的 pickle 去反序列化。
(2)pickle 能执行任意对象,如果权重文件被篡改,加载时可能执行恶意代码。

2.什么时候会出问题

目前这只是一个 警告,程序还能跑。但在未来版本中,PyTorch 将把默认值改成:
这会带来两种情况:
(1)推理/迁移学习:只需要权重张量 (state_dict),完全不受影响,反而更安全。
(2)继续训练:如果 checkpoint 里还有优化器、调度器、自定义对象,那么未来就会加载失败,除非你显式写 weights_only=False。

3.解决方案

显式开启 weights_only=True(立刻可用)
ckpt = torch.load(path, map_location="cpu", weights_only=True)

您可能感兴趣的与本文相关的镜像

PyTorch 2.5

PyTorch 2.5

PyTorch
Cuda

PyTorch 是一个开源的 Python 机器学习库,基于 Torch 库,底层由 C++ 实现,应用于人工智能领域,如计算机视觉和自然语言处理

PyTorch 中加载模型权重时,可能会遇到多种问题,尤其是在模型结构或训练环境发生变化的情况下。以下是一些常见的问题及其解决方案: ### 加载模型权重的基本方法 当使用 `torch.load()` 加载模型权重时,通常需要确保模型的结构与保存权重时的结构一致。可以通过以下方式加载模型权重: ```python model = TheModelClass() # 确保模型类已经定义 model.load_state_dict(torch.load(PATH)) model.eval() # 设置模型为评估模式 ``` 如果模型是在多 GPU 环境下训练并使用 `torch.nn.DataParallel` 进行包装,则保存的权重会包含 `module` 前缀。为了在单个 GPU 或 CPU 上加载这些权重,需要对权重文件进行调整,去除 `module` 前缀: ```python def remove_module_prefix(state_dict): from collections import OrderedDict new_state_dict = OrderedDict() for k, v in state_dict.items(): name = k[7:] # 去除 'module.' 前缀 new_state_dict[name] = v return new_state_dict # 加载权重 checkpoint = torch.load(PATH, map_location=device) model.load_state_dict(remove_module_prefix(checkpoint)) ``` ### 模型训练状态恢复 在恢复模型训练状态时,除了加载模型权重外,还需要加载优化器状态和其他训练相关信息。通常,这些信息会在保存模型时一并保存: ```python checkpoint = torch.load(PATH) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) epoch = checkpoint['epoch'] loss = checkpoint['loss'] model.train() # 设置模型为训练模式 ``` ### 处理模型结构不一致的情况 如果模型结构发生了变化,导致无法直接加载权重,可以手动匹配权重。例如,如果新模型和旧模型的部分层名称不同,可以手动将旧模型权重赋值给新模型的相应层: ```python model_dict = model.state_dict() pretrained_dict = torch.load(PATH) # 过滤出模型中需要的层 pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and model_dict[k].size() == v.size()} # 更新模型权重 model_dict.update(pretrained_dict) model.load_state_dict(model_dict) ``` ### 自定义加载逻辑 在某些情况下,可能需要自定义加载逻辑来处理特定的需求。例如,如果模型中某些层的权重形状不匹配,可以选择跳过这些层的加载: ```python model_dict = model.state_dict() pretrained_dict = torch.load(PATH) for k, v in pretrained_dict.items(): if k in model_dict and model_dict[k].size() == v.size(): model_dict[k] = v else: print(f"Skipping {k} due to size mismatch.") model.load_state_dict(model_dict) ``` ###
评论
成就一亿技术人!
拼手气红包6.0元
还能输入1000个字符
 
红包 添加红包
表情包 插入表情
 条评论被折叠 查看
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值