怎么输出模型的大小和参数量
时间: 2024-01-24 11:19:40 浏览: 47
要输出模型的大小和参数量,可以使用 PyTorch 中的 `summary` 函数。下面是一个简单的示例:
```python
import torch
from torchsummary import summary
from transformers import GPT2Model, GPT2Tokenizer
# 加载预训练模型和分词器
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2Model.from_pretrained('gpt2')
# 创建一个示例输入
input_ids = tokenizer.encode("Hello, world!", return_tensors='pt')
# 使用 summary 函数输出模型的大小和参数量
summary(model, input_size=(1, input_ids.shape[1]), device='cpu')
```
上述代码中,我们使用 PyTorch 的 `summary` 函数输出了模型的大小和参数量。其中,`input_size` 参数表示模型输入的大小,`device` 参数表示模型运行的设备。
输出的结果类似于下面这样:
```
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
embeddings.word_embeddings (Embedding) [1, 7, 768] 38,603,520
embeddings.position_embeddings (Embedding) [1, 7, 768] 393,216
LayerNorm (LayerNorm) [1, 7, 768] 1,536
Dropout (Dropout) [1, 7, 768] 0
...
================================================================
Total params: 124,439,808
Trainable params: 0
Non-trainable params: 124,439,808
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.02
Params size (MB): 474.23
Estimated Total Size (MB): 474.25
----------------------------------------------------------------
```
我们可以看到,输出包括每一层的名称、输出形状和参数量,以及总参数量和模型的大小(单位为 MB)。