加载 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