如何估算基于ENRON数据集训练的DistilBERT模型大小及相关疑问

用户问题 #

我需要获取基于ENRON数据集训练的DistilBERT模型的大小，打算通过各层参数数量、权重与偏置来近似估算，但不清楚具体方法，该如何计算？另外，调用model.summary()时，输出形状列显示为“multiple”，这是什么原因？

相关代码与信息 #

模型实现代码

import datetime
import pandas as pd
import tensorflow as tf
import transformers
from transformers import DistilBertTokenizer
from transformers import TFAutoModelForSequenceClassification
pd.set_option('display.max_colwidth', None)
MODEL_NAME = 'distilbert-base-uncased'

# 加载并处理数据
df = pd.read_csv("enron6dataset.csv", error_bad_lines=False)
train_df = df.sample(frac=0.8, random_state=25)
test_df = df.drop(train_df.index)
X_train = train_df.email
X_test = test_df.email
y_train = train_df.label
y_test = test_df.label

y_train_encoded = tf.one_hot(y_train.values, 2)
y_test_encoded = tf.one_hot(y_test.values, 2)

# 分词与模型编译
tokenizer = DistilBertTokenizer.from_pretrained(MODEL_NAME)

train_encodings = tokenizer(list(X_train.values),
                            truncation=True, 
                            padding=True)
test_encodings = tokenizer(list(X_test.values),
                           truncation=True, 
                           padding=True)

train_dataset = tf.data.Dataset.from_tensor_slices((dict(train_encodings),list(y_train_encoded)))

test_dataset = tf.data.Dataset.from_tensor_slices((dict(test_encodings),list(y_test_encoded)))
test_dataset2 = test_dataset.shuffle(buffer_size=1024).take(1000).batch(16)

model = TFAutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

optimizerr = tf.keras.optimizers.Adam(learning_rate=5e-5)

model.compile(optimizer=optimizerr,
              loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
              metrics=['accuracy'])

# 微调模型
model.fit(train_dataset.shuffle(len(X_train)).batch(16),
          epochs=2,
          batch_size=16,callbacks=[tensorboard_callback])
# 测试精度
results = model.evaluate(test_dataset2)
print("test loss, test acc:", results)

模型测试精度

63/63 [==============================] - 2s 32ms/step - loss: 0.0135 - accuracy: 0.9950
test loss, test acc: [0.01349691767245531, 0.9950000047683716]

model.summary()输出

Model: "tf_distil_bert_for_sequence_classification"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
distilbert (TFDistilBertMain multiple                  66362880  
_________________________________________________________________
pre_classifier (Dense)       multiple                  590592    
_________________________________________________________________
classifier (Dense)           multiple                  1538      
_________________________________________________________________
dropout_19 (Dropout)         multiple                  0         
=================================================================
Total params: 66,955,010
Trainable params: 66,955,010
Non-trainable params: 0
_________________________________________________________________

模型层信息

[<transformers.models.distilbert.modeling_tf_distilbert.TFDistilBertMainLayer at 0x2b268cbbfb20>,
 <tensorflow.python.keras.layers.core.Dense at 0x2b268cbdd670>,
 <tensorflow.python.keras.layers.core.Dense at 0x2b268cbddfa0>,
 <tensorflow.python.keras.layers.core.Dropout at 0x2b268cbdf730>]

问题解答 #

1. 如何计算模型大小？ #

模型大小可以通过参数数量×单个参数占用字节数来近似估算，具体步骤如下：

• 第一步：获取总参数数。从model.summary()中可以看到总参数为66,955,010。
• 第二步：确定参数存储格式。默认情况下，TensorFlow中模型参数以FP32（单精度浮点数）存储，每个参数占用4字节；如果是FP16（半精度）则占用2字节。
• 第三步：计算总字节数并转换单位：
  • FP32格式：66,955,010 × 4 = 267,820,040 字节，换算为MB的话，除以1024×1024（约1,048,576），得到≈255.4MB。
  • 如果是FP16格式，大小则是≈127.7MB。
• 补充：实际磁盘存储大小会和这个估算值接近，但可能因为模型保存时的优化（如权重压缩、量化）或附带的配置文件、分词器文件略有差异。你也可以直接保存模型后查看文件大小，比如用model.save_pretrained("./enron_distilbert")，然后查看文件夹的占用空间。

2. 为什么输出形状列显示“multiple”？ #

出现“multiple”的原因有两个：

• 首先，distilbert层是嵌套的复杂结构（TFDistilBertMainLayer），它本身会输出多个张量，比如Transformer最后一层的隐藏状态、池化后的序列表示等，并非单一输出形状，Keras无法在summary中用单个形状表示，因此标记为“multiple”。
• 其次，Hugging Face的TensorFlow模型部分层的输出形状是动态的，会根据输入序列长度变化，或者层设计为多输出模式，这也会导致summary无法显示固定的单一形状。

如果需要查看具体的输出形状，可以手动传入一个测试输入：

# 用tokenizer编码一个测试句子
test_input = tokenizer("This is a test email", return_tensors="tf")
# 获取模型输出
output = model(test_input)
# 打印输出形状
print(output.logits.shape)

内容的提问来源于stack exchange，提问作者Sri991