Bootstrap

python计算stable-diffusion-1.5模型参数量以及该模型每一层网络的参数量【其他LLM模型也有参考意义】

最近在计算stable-diffusion-1.5模型参数量上花了点心思,总结了一些方法,一起学习:

stable-diffusion-1.5模型结构

首先stable-diffusion-1.5模型主要有三个关键组件(text_encoder,unet,vae),关于stable-diffusion-1.5模型,有个大佬总结的帖子十分到位,推荐大家学习,链接:
Stable Diffusion1.5网络结构-超详细原创

计算模型的整体参数量

那么对于这种大模型,我们又该怎么计算其参数量呢?
我的思路是直接读取文件,分别计算三个关键组件(text_encoder,unet,vae)的参数量,再相加,直接上代码:

from diffusers import StableDiffusionPipeline
import torch

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def print_layer_info(model):
    for name, layer in model.named_modules():
        if hasattr(layer, 'parameters'):
            layer_params = sum(p.numel() for p in layer.parameters() if p.requires_grad)
            print(f"Layer: {name}, Parameters: {layer_params}")

pipe = StableDiffusionPipeline.from_pretrained(
    "xxxxxxx/stable - diffusion - v1 - 5",
    torch_dtype=torch.float16
)
pipe = pipe.to("cuda")

# 分别计算各个组件的参数
num_params_text_encoder = count_parameters(pipe.text_encoder)
num_params_unet = count_parameters(pipe.unet)
num_params_vae = count_parameters(pipe.vae)

total_num_params = num_params_text_encoder + num_params_unet + num_params_vae
print(f"Number of parameters in Text Encoder: {num_params_text_encoder}")
print(f"Number of parameters in UNet: {num_params_unet}")
print(f"Number of parameters in VAE: {num_params_vae}")
print(f"Total number of parameters: {total_num_params}")

# 打印每个部件每一层的信息和参数量信息
print("Text Encoder Layers Information:")
print_layer_info(pipe.text_encoder)

print("UNet Layers Information:")
print_layer_info(pipe.unet)

print("VAE Layers Information:")
print_layer_info(pipe.vae)

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]
image.save("astronaut_rides_horse.png")

注意:需要将"xxxxxxx/stable - diffusion - v1 - 5",这个地址改成你自己保存源码的地址。

上面代码的输出内容如下:
在这里插入图片描述

精细的输出每一部分的每一层的参数信息

又有一个问题,如果我想要看到每一个组件中的每一层的参数量,又该怎么做?直接上代码:

from diffusers import StableDiffusionPipeline
import torch

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def print_layer_info(model):
    for name, layer in model.named_modules():
        if hasattr(layer, 'parameters'):
            layer_params = sum(p.numel() for p in layer.parameters() if p.requires_grad)
            print(f"Layer: {name}, Parameters: {layer_params}")


pipe = StableDiffusionPipeline.from_pretrained(
    "xxxxxxxxxx/stable - diffusion - v1 - 5",
    torch_dtype=torch.float16
)
pipe = pipe.to("cuda")

# 分别计算各个组件的参数
num_params_text_encoder = count_parameters(pipe.text_encoder)
num_params_unet = count_parameters(pipe.unet)
num_params_vae = count_parameters(pipe.vae)

total_num_params = num_params_text_encoder + num_params_unet + num_params_vae
print(f"Number of parameters in Text Encoder: {num_params_text_encoder}")
print(f"Number of parameters in UNet: {num_params_unet}")
print(f"Number of parameters in VAE: {num_params_vae}")
print(f"Total number of parameters: {total_num_params}")

# 打开文件用于保存输出信息
with open('model_info.txt', 'w') as f:
    # 将总参数量信息写入文件
    f.write(f"Number of parameters in Text Encoder: {num_params_text_encoder}\n")
    f.write(f"Number of parameters in UNet: {num_params_unet}\n")
    f.write(f"Number of parameters in VAE: {num_params_vae}\n")
    f.write(f"Total number of parameters: {total_num_params}\n")

    # 打印并保存每个部件每一层的信息和参数量信息
    f.write("Text Encoder Layers Information:\n")
    for name, layer in pipe.text_encoder.named_modules():
        if hasattr(layer, 'parameters'):
            layer_params = sum(p.numel() for p in layer.parameters() if p.requires_grad)
            f.write(f"Layer: {name}, Parameters: {layer_params}\n")
    f.write("\n")

    f.write("UNet Layers Information:\n")
    for name, layer in pipe.unet.named_modules():
        if hasattr(layer, 'parameters'):
            layer_params = sum(p.numel() for p in layer.parameters() if p.requires_grad)
            f.write(f"Layer: {name}, Parameters: {layer_params}\n")
    f.write("\n")

    f.write("VAE Layers Information:\n")
    for name, layer in pipe.vae.named_modules():
        if hasattr(layer, 'parameters'):
            layer_params = sum(p.numel() for p in layer.parameters() if p.requires_grad)
            f.write(f"Layer: {name}, Parameters: {layer_params}\n")

上述代码将该过程的打印信息都保存在了当前目录的model_info.txt文件下。这样就会看到每一层的参数信息了。在这里插入图片描述
以上就是我的两种方法,对于大部分LLM模型,这种读取参数的操作都是适用的,如果有伙伴想尝试读取其他模型的参数,也可以借鉴。

;