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 torchdef 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 torchdef 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")