加载 safetensors 格式模型
torch 加载 safetensors 格式的模型
文章共228字 · 阅读需要大约1分钟
一键AI生成摘要,助你高效阅读
问答
·
# coding:utf-8
# @Email: wangguisen@donews.com
# @Time: 2023/3/22 17:39
# @File: ddd.py
'''
https://github.com/huggingface/safetensors
https://huggingface.co/docs/safetensors/index
pip install safetensors
torch >= 2.0
'''
import torch
from safetensors.torch import load_file, save_file
from safetensors import safe_open
model_path = 'nutriasOtters_10.safetensors'
tensors = {}
with safe_open(model_path, framework="pt", device='cpu') as f:
for k in f.keys():
tensors[k] = f.get_tensor(k)
print(tensors)
# model.load_state_dict
更多推荐
已为社区贡献2条内容
所有评论(0)