如何自定义预训练Transformer的编码器数量？以BART文本摘要为例

我正在使用Python的transformers模块实现预训练Transformer模型以完成文本摘要任务，希望对比不同编码器数量下微调后BART模型的性能。我的疑问是：如何自定义编码器数量？默认BART拥有12个编码器，若只想保留前6个该如何操作？我找到了BART的配置文档，但不知道如何适配现有代码，现有代码如下：

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer, DataCollatorForSeq2Seq

tokenizer = AutoTokenizer.from_pretrained(model_checkpoints)
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoints)
collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)

# preprocessing step omitted
# tokenized_data = preprocessed data

args = transformers.Seq2SeqTrainingArguments(
    'conversation-summ',
    evaluation_strategy='epoch',
    learning_rate=2e-5,
    per_device_train_batch_size=1,
    per_device_eval_batch_size= 1,
    gradient_accumulation_steps=2,
    weight_decay=0.01,
    save_total_limit=2,
    num_train_epochs=3,
    predict_with_generate=True,
    eval_accumulation_steps=1,
    fp16=True
    )

trainer = transformers.Seq2SeqTrainer(
    model, 
    args,
    train_dataset=tokenized_data['train'],
    eval_dataset=tokenized_data['validation'],
    data_collator=collator,
    tokenizer=tokenizer,
    compute_metrics=compute_rouge
)

trainer.train()

核心逻辑 #

BART的编码器层数由配置文件中的encoder_layers参数控制，但直接修改配置后加载模型会随机初始化新层，丢失预训练权重。正确做法是加载完整预训练模型后，保留前N层编码器，同时修改配置让模型适配该层数，这样既能复用预训练知识，又能实现层数自定义。

具体步骤 #

1. 修改配置并复用预训练权重 #

替换原代码中模型加载的部分，改为以下代码：

from transformers import AutoTokenizer, AutoConfig, AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer, DataCollatorForSeq2Seq

model_checkpoints = "facebook/bart-base"  # 可替换为你使用的自定义checkpoint
tokenizer = AutoTokenizer.from_pretrained(model_checkpoints)

# 步骤1：加载原模型配置并修改编码器层数为6
config = AutoConfig.from_pretrained(model_checkpoints)
config.encoder_layers = 6

# 步骤2：加载完整的预训练BART模型
full_pretrained_model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoints)

# 步骤3：用修改后的配置初始化新模型，然后替换编码器层为原模型的前6层
model = AutoModelForSeq2SeqLM.from_config(config)
# 替换编码器的Transformer层：仅保留原模型的前6层
model.model.encoder.layers = full_pretrained_model.model.encoder.layers[:6]
# 保留预训练的嵌入层和归一化层（这些是模型的基础组件，不能随机初始化）
model.model.encoder.embed_tokens = full_pretrained_model.model.encoder.embed_tokens
model.model.encoder.embed_positions = full_pretrained_model.model.encoder.embed_positions
model.model.encoder.layernorm_embedding = full_pretrained_model.model.encoder.layernorm_embedding
model.model.encoder.final_layer_norm = full_pretrained_model.model.encoder.final_layer_norm

2. 验证修改结果 #

添加一行代码确认编码器层数是否正确修改：

print(f"修改后的编码器层数：{len(model.model.encoder.layers)}")  # 输出应为6

3. 后续代码复用原逻辑 #

DataCollator、训练参数、Trainer的定义和原代码完全一致，直接保留即可：

collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)

# preprocessing step omitted
# tokenized_data = preprocessed data

args = Seq2SeqTrainingArguments(
    'conversation-summ',
    evaluation_strategy='epoch',
    learning_rate=2e-5,
    per_device_train_batch_size=1,
    per_device_eval_batch_size= 1,
    gradient_accumulation_steps=2,
    weight_decay=0.01,
    save_total_limit=2,
    num_train_epochs=3,
    predict_with_generate=True,
    eval_accumulation_steps=1,
    fp16=True
    )

trainer = Seq2SeqTrainer(
    model, 
    args,
    train_dataset=tokenized_data['train'],
    eval_dataset=tokenized_data['validation'],
    data_collator=collator,
    tokenizer=tokenizer,
    compute_metrics=compute_rouge
)

trainer.train()

关键细节说明 #

• 不能直接用修改后的配置加载预训练模型：如果直接调用AutoModelForSeq2SeqLM.from_pretrained(model_checkpoints, config=config)，模型会随机初始化缺少的编码器层，导致预训练权重失效，模型性能骤降。
• 保留基础组件：编码器的嵌入层、位置嵌入层和归一化层是整个模型的输入输出基础，必须复用预训练权重，否则输入的文本表示会偏离预训练分布。
• 对比不同层数：只需修改config.encoder_layers的值（比如3、6、9），重复上述步骤即可快速生成不同编码器层数的BART模型，用于性能对比实验。

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