转 https://blog.csdn.net/wzk4869/article/details/130668642
显卡:GTX 3060
safetensors
safetensors 是一种用于安全存储张量(与 pickle 相反)的新型简单格式,并且仍然很快(零拷贝)。
1.1 pip 安装
pip install safetensors
conda install -c huggingface safetensors
1.2 加载张量
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)
# 仅加载部分张量(在多个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]
1.3 保存张量
import torch
from safetensors.torch import save_file
tensors = {
"embedding": torch.zeros((2, 2)),
"attention": torch.zeros((2, 3))
}
save_file(tensors, "model.safetensors")
1.4 速度比较
safetensors 真的很快。让我们通过加载 gpt2 权重将其进行比较。要运行 GPU 基准测试,请确保您的机器具有 GPU,或者您已选择是否使用的是 Google Colab。
在开始之前,请确保已安装所有必要的库:
pip install safetensors huggingface_hub torch
让我们从导入所有将使用的包开始:
import os
import datetime
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import torch
# Download safetensors & torch weights for gpt2
sf_filename = hf_hub_download("gpt2", filename="model.safetensors")
pt_filename = hf_hub_download("gpt2", filename="pytorch_model.bin")
# CPU 基准测试
start_st = datetime.datetime.now()
weights = load_file(sf_filename, device="cpu")
load_time_st = datetime.datetime.now() - start_st
print(f"Loaded safetensors {load_time_st}")
# 输出结果为:Loaded safetensors 0:00:00.026842
start_pt = datetime.datetime.now()
weights = torch.load(pt_filename, map_location="cpu")
load_time_pt = datetime.datetime.now() - start_pt
print(f"Loaded pytorch {load_time_pt}")
# 输出结果为:Loaded pytorch 0:00:00.182266
print(f"on CPU, safetensors is faster than pytorch by: {load_time_pt/load_time_st:.1f} X")
# 输出结果为:on CPU, safetensors is faster than pytorch by: 6.8 X
这种加速是由于该库通过直接映射文件来避免不必要的副本。实际上可以在 torch 上完成。
# GPU 基准测试
os.environ["SAFETENSORS_FAST_GPU"] = "1"
torch.zeros((2, 2)).cuda()
start_st = datetime.datetime.now()
weights = load_file(sf_filename, device="cuda:0")
load_time_st = datetime.datetime.now() - start_st
print(f"Loaded safetensors {load_time_st}")
start_pt = datetime.datetime.now()
weights = torch.load(pt_filename, map_location="cuda:0")
load_time_pt = datetime.datetime.now() - start_pt
print(f"Loaded pytorch {load_time_pt}")
print(f"on GPU, safetensors is faster than pytorch by: {load_time_pt/load_time_st:.1f} X")
# 输出结果为:
# Loaded safetensors 0:00:00.497415
# Loaded pytorch 0:00:00.250602
# on GPU, safetensors is faster than pytorch by: 0.5 X
加速有效是因为此库能够跳过不必要的 CPU 分配。不幸的是,据我们所知,它无法在纯 pytorch 中复制。该库的工作原理是内存映射文件,使用 pytorch 创建空张量,并直接调用以直接在 GPU 上移动张量。
==========================================================
实际操作中,不知道为什么上面的代码是跑不通的
打开safetensors格式文件要换safetensors库的safe_open
from safetensors import safe_open
tensors1 = {}
with safe_open(filex, framework="pt", device="cuda:0") as f:
for k in f.keys():
tensors1[k] = f.get_tensor(k)
safe_open下面也没有modules属性
safetensors库运行时发现modules()函数不存在
torch.nn.Module的children()函数也不存在,named_children()也不存在
pytorch版本 pytorch-2.1.0-py3.9_cpu_0
具体统计卷积层数量,池化层数量,全连接层数量的代码(会报错)
import torch
import torch.nn as nn
# 加载模型
filex = 'D:/Program Files/sdxl-webui-1.0/sdxl-webui-1.0/models/Stable-diffusion/sd_xl_base_1.0.safetensors'
filexs = 'C:/Users/admin/Desktop/model.pth'
# torch.save(filex.state_dict(), filexs)
# model.load_state_dict(torch.load(filexs))
# model = torch.load('model.pth')
# 统计各层数量
conv_count = 0
pool_count = 0
fc_count = 0
# 遍历模型各层
for layer in model.modules():
if isinstance(layer, nn.Conv2d):
conv_count += 1
elif isinstance(layer, nn.MaxPool2d):
pool_count += 1
elif isinstance(layer, nn.Linear):
fc_count += 1
# 输出各层数量
print("卷积层数量:", conv_count)
print("池化层数量:", pool_count)
print("全连接层数量:", fc_count)
# ===================================================
import torch.nn as nn
from safetensors import safe_open
# 加载模型
filex = 'D:/Program Files/sdxl-webui-1.0/sdxl-webui-1.0/models/Stable-diffusion/sd_xl_base_1.0.safetensors'
filexs = 'C:/Users/admin/Desktop/model2.safetensors'
model = safe_open(filex, framework="pt", device="cuda:0") # 加载模型的完整路径
def get_modules(module):
modules = [module]
for name, child in module.nn.named_children():
modules += get_modules(child)
return modules
# 统计各层数量
conv_count = 0
pool_count = 0
fc_count = 0
modules = get_modules(model)
for module in modules:
if isinstance(module, nn.Conv2d):
conv_count += 1
elif isinstance(module, nn.MaxPool2d):
pool_count += 1
elif isinstance(module, nn.Linear):
fc_count += 1
# 输出各层数量
print("卷积层数量:", conv_count)
print("池化层数量:", pool_count)
print("全连接层数量:", fc_count)
# ===================================================
import torch.nn as nn
from safetensors import safe_open
# 加载模型
filex = 'D:/Program Files/sdxl-webui-1.0/sdxl-webui-1.0/models/Stable-diffusion/sd_xl_base_1.0.safetensors'
filexs = 'C:/Users/admin/Desktop/model2.safetensors'
model = safe_open(filex, framework="pt", device="cuda:0") # 加载模型的完整路径
# 保存模型的状态字典
# torch.save(model.state_dict(), filexs)
# 重新加载模型的状态字典
# model_state_dict = torch.load(filexs)
def get_modules(module):
modules = [module]
for name, child in module.nn.named_children():
modules += get_modules(child)
return modules
# 统计各层数量
conv_count = 0
pool_count = 0
fc_count = 0
modules = get_modules(model)
for module in modules:
if isinstance(module, nn.Conv2d):
conv_count += 1
elif isinstance(module, nn.MaxPool2d):
pool_count += 1
elif isinstance(module, nn.Linear):
fc_count += 1
# 输出各层数量
print("卷积层数量:", conv_count)
print("池化层数量:", pool_count)
print("全连接层数量:", fc_count)
暂时还没找到原因