Hugging Face, EleutherAI, StabilityAI 用的多

介紹

文件形式

• header，體現(xiàn)其特性。如果強(qiáng)行將pickle或者空軟連接 打開(kāi)，會(huì)出現(xiàn)報(bào)錯(cuò)。解決詳見(jiàn)：debug 連接到其他教程
• 結(jié)構(gòu)和參數(shù)
  

安裝

with pip:Copied
pip install safetensors
with conda:Copied
conda install -c huggingface safetensors

Usage

文檔: https://huggingface.co/docs/safetensors/index
github: https://github.com/huggingface/safetensors

測(cè)試安裝

import torch
from safetensors import safe_open
from safetensors.torch import save_filetensors = {"weight1": torch.zeros((1024, 1024)),"weight2": torch.zeros((1024, 1024))
}
save_file(tensors, "model.safetensors")tensors = {}
with safe_open("model.safetensors", framework="pt", device="cpu") as f:for key in f.keys():tensors[key] = f.get_tensor(key)

加載

文檔 https://huggingface.co/docs/diffusers/using-diffusers/using_safetensors

from diffusers import StableDiffusionPipelinepipeline = StableDiffusionPipeline.from_single_file("https://huggingface.co/WarriorMama777/OrangeMixs/blob/main/Models/AbyssOrangeMix/AbyssOrangeMix.safetensors"
)

Load tensors

from safetensors import safe_opentensors = {}
with safe_open("model.safetensors", framework="pt", device=0) as f:for k in f.keys():tensors[k] = f.get_tensor(k)

# Loading only part of the tensors (interesting when running on multiple GPU)from safetensors import safe_opentensors = {}
with safe_open("model.safetensors", framework="pt", device=0) as f:tensor_slice = f.get_slice("embedding")vocab_size, hidden_dim = tensor_slice.get_shape()tensor = tensor_slice[:, :hidden_dim]

保存

import torch
from safetensors.torch import save_filetensors = {"embedding": torch.zeros((2, 2)),"attention": torch.zeros((2, 3))
}
save_file(tensors, "model.safetensors")

轉(zhuǎn)換到safetensor

• 在線，利用hugging face

The easiest way to convert your model weights is to use the Convert Space, given your model weights are already stored on the Hub. The Convert Space downloads the pickled weights, converts them, and opens a Pull Request to upload the newly converted .safetensors file to your repository.

• 本地 運(yùn)行
  see 轉(zhuǎn)換代碼 convert.py

# 主函數(shù)
def convert_file(pt_filename: str,sf_filename: str,
):loaded = torch.load(pt_filename, map_location="cpu")if "state_dict" in loaded:loaded = loaded["state_dict"]shared = shared_pointers(loaded)for shared_weights in shared:for name in shared_weights[1:]:loaded.pop(name)# For tensors to be contiguousloaded = {k: v.contiguous() for k, v in loaded.items()}dirname = os.path.dirname(sf_filename)os.makedirs(dirname, exist_ok=True)save_file(loaded, sf_filename, metadata={"format": "pt"})check_file_size(sf_filename, pt_filename)reloaded = load_file(sf_filename)for k in loaded:pt_tensor = loaded[k]sf_tensor = reloaded[k]if not torch.equal(pt_tensor, sf_tensor):raise RuntimeError(f"The output tensors do not match for key {k}")

例子

解析

import requests # pip install requests
import structdef parse_single_file(url):# Fetch the first 8 bytes of the fileheaders = {'Range': 'bytes=0-7'}response = requests.get(url, headers=headers)# Interpret the bytes as a little-endian unsigned 64-bit integerlength_of_header = struct.unpack('<Q', response.content)[0]# Fetch length_of_header bytes starting from the 9th byteheaders = {'Range': f'bytes=8-{7 + length_of_header}'}response = requests.get(url, headers=headers)# Interpret the response as a JSON objectheader = response.json()return headerurl = "https://huggingface.co/gpt2/resolve/main/model.safetensors"
header = parse_single_file(url)print(header)