快速入门Safetensors
快速入门Safetensors
-
- 什么是Safetensors
- 架构
- 常用操作
- 速度对比
- 彩蛋
Safetensors官方网址
什么是Safetensors
Safetensors是一种新的简单格式，用于安全存储张量(与pickle相反)，而且速度仍然很快(零拷贝)。
架构
常用操作
# pip install safetensors
# Load
from safetensors import safe_open
tensors = {}
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_open
tensors = {}
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_file
tensors = {
"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 json
with 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] = v
if isinstance(v, str) and v[0:1] == '{':
try:
res[k] = json.loads(v)
except Exception:
pass
return 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"
}