使用HuggingFace库在Pytorch中训练n%的最后一层BERT (训练12个中的最后5个BERTLAYER )

pip install transformers torch

from transformers import BertForSequenceClassification

model = BertForSequenceClassification.from_pretrained('bert-base-uncased')

# 冻结除最后5层之外的所有层
for name, param in model.bert.named_parameters():
    if name.split('.')[2] not in ['9', '10', '11']:  # 这里的数字代表要训练的层（从0开始计数）
        param.requires_grad = False

from torch.optim import AdamW

# 只为需要梯度的参数设置优化器
optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-5)

from torch.utils.data import DataLoader

# 假设你已经有了一个PyTorch数据集 `dataset`
train_loader = DataLoader(dataset, batch_size=16, shuffle=True)

model.train()
for epoch in range(num_epochs):
    for batch in train_loader:
        inputs = batch['input_ids']
        labels = batch['labels']
        outputs = model(inputs, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

model.eval()
# 进行评估