加载 safetensors 格式模型_torch 加载safetensors

# 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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26