快速入门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"
    }