Safetensors官方网址

什么是Safetensors

Safetensors是一种新的简单格式，用于安全存储张量(与pickle相反)，而且速度仍然很快(零拷贝)。

架构

常用操作

# pip install safetensors# Load
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]# save
import torch
from safetensors.torch import save_filetensors = {"embedding": torch.zeros((2, 2)),"attention": torch.zeros((2, 3))
}
save_file(tensors, "model.safetensors")

速度对比

• cpu

Loaded safetensors 0:00:00.004015
Loaded pytorch 0:00:00.307460
on CPU, safetensors is faster than pytorch by: 76.6 X

• gpu

Loaded safetensors 0:00:00.165206
Loaded pytorch 0:00:00.353889
on GPU, safetensors is faster than pytorch by: 2.1 X

彩蛋

Lora的Safetensors格式存储有训练时的元数据信息

代码来自sd-web-ui项目源码

def read_metadata_from_safetensors(filename):import jsonwith open(filename, mode="rb") as file:metadata_len = file.read(8)metadata_len = int.from_bytes(metadata_len, "little")json_start = file.read(2)assert metadata_len > 2 and json_start in (b'{"', b"{'"), f"{filename} is not a safetensors file"json_data = json_start + file.read(metadata_len-2)json_obj = json.loads(json_data)res = {}for k, v in json_obj.get("__metadata__", {}).items():res[k] = vif isinstance(v, str) and v[0:1] == '{':try:res[k] = json.loads(v)except Exception:passreturn res
import sys
print(read_metadata_from_safetensors(sys.argv[1]))

{"ss_sd_model_name": "dream-shaper-5.ckpt","ss_resolution": "(512, 512)","ss_clip_skip": "2","ss_num_train_images": "5000","ss_tag_frequency": {"100_coloredic0n": {"coloredic0n icon a white ambulance with red stripe on it": 1,"coloredic0n icon a brown backpack": 1,"coloredic0n icon a banana peel": 1,"coloredic0n icon a battery with a lightning on it": 1,"coloredic0n icon a bed with a pillow on it": 1,"coloredic0n icon a red book with the words on it": 1,"coloredic0n icon a cabbage with green leaves": 1,"coloredic0n icon a cactus plant in a pot": 1,"coloredic0n icon a cassette tape": 1,"coloredic0n icon a  champagne bottle": 1,"coloredic0n icon a piece of cheese": 1,"coloredic0n icon a squirrel sitting with a nut in its hands": 1,"coloredic0n icon a cookie with a chocolate chip on top": 1,"coloredic0n icon a red price percent symbol on it": 1,"coloredic0n icon a fence": 1,"coloredic0n icon a red fire hydrant": 1,"coloredic0n icon a woman with long hair smiling and wearing a pink shirt": 1,"coloredic0n icon a pair of glasses": 1,"coloredic0n icon a pink jellyfish floating in the air": 1,"coloredic0n icon a ladybug with a black and red pattern on it": 1,"coloredic0n icon a green lawnmower": 1,"coloredic0n icon a carton of milk with a cap on top": 1,"coloredic0n icon a bowl of noodles with peas and tomatoes": 1,"coloredic0n icon a notebook with a pen on top of it": 1,"coloredic0n icon a onion": 1,"coloredic0n icon a padlock": 1,"coloredic0n icon a panda face": 1,"coloredic0n icon a blue pen": 1,"coloredic0n icon a pencil": 1,"coloredic0n icon a black power plug": 1,"coloredic0n icon a rat with a long tail": 1,"coloredic0n icon a raven": 1,"coloredic0n icon a rifle with a bullets": 1,"coloredic0n icon a red and white rocket ship": 1,"coloredic0n icon a red rose with green leaves": 1,"coloredic0n icon a yellow rubber duck": 1,"coloredic0n icon a sneaker with a red and white sole": 1,"coloredic0n icon a green snail with a yellow shell": 1,"coloredic0n icon a pair of socks with a red and white design": 1,"coloredic0n icon a brown couch with pillows and a pillow on it": 1,"coloredic0n icon a grey sofa with pillows and a pillow on it": 1,"coloredic0n icon a sunflower": 1,"coloredic0n icon a yellow ticket with a star on it": 1,"coloredic0n icon a toilet seat with a lid up": 1,"coloredic0n icon a wood log": 1,"coloredic0n icon a red and white target with an arrow in the center": 1,"coloredic0n icon a bucket filled with water": 1,"coloredic0n icon a calendar": 1,"coloredic0n icon a red pencil sharper with a metal blade": 1,"coloredic0n icon a red and white and red megaphone": 1}},"ss_batch_size_per_device": "2","ss_bucket_info": "null","ss_bucket_no_upscale": "False","ss_cache_latents": "True","ss_caption_dropout_every_n_epochs": "0","ss_caption_dropout_rate": "0.0","ss_caption_tag_dropout_rate": "0.0","ss_color_aug": "False","ss_dataset_dirs": {"100_coloredic0n": {"n_repeats": 100,"img_count": 50}},"ss_enable_bucket": "False","ss_epoch": "1","ss_face_crop_aug_range": "None","ss_flip_aug": "False","ss_full_fp16": "False","ss_gradient_accumulation_steps": "1","ss_gradient_checkpointing": "False","ss_keep_tokens": "0","ss_learning_rate": "0.0001","ss_lowram": "False","ss_lr_scheduler": "constant","ss_lr_warmup_steps": "0","ss_max_bucket_reso": "None","ss_max_grad_norm": "1.0","ss_max_token_length": "None","ss_max_train_steps": "2500","ss_min_bucket_reso": "None","ss_min_snr_gamma": "None","ss_mixed_precision": "bf16","ss_network_alpha": "Dynamic","ss_network_dim": "Dynamic","ss_network_module": "networks.lora","ss_new_sd_model_hash": "ffad8f3da766de2a56343e388df43941ee27143a03c596e0b8f06976d4fa52e9","ss_noise_offset": "None","ss_num_batches_per_epoch": "2500","ss_num_epochs": "1","ss_num_reg_images": "0","ss_optimizer": "bitsandbytes.optim.adamw.AdamW8bit","ss_output_name": "Colored_Icons","ss_prior_loss_weight": "1.0","ss_random_crop": "False","ss_reg_dataset_dirs": {},"ss_sd_model_hash": "02a6900f","ss_sd_scripts_commit_hash": "9533285e998c74b0e18d114aea71f7143c8f3fa9","ss_seed": "1234","ss_session_id": "2548637574","ss_shuffle_caption": "False","ss_text_encoder_lr": "5e-05","ss_total_batch_size": "2","ss_training_comment": "Dynamic resize with sv_fro: 0.9 from 128; None","ss_training_finished_at": "1683314416.566944","ss_training_started_at": "1683312438.3090656","ss_unet_lr": "0.0001","ss_v2": "False","sshs_legacy_hash": "a45a7daa","sshs_model_hash": "6a455e4a89e4c0cbb5985bf8ebb96397e61bb4b8431b121281d6ec796daecae1"}