PyTorch下BERT模型训练损失与准确率曲线绘制方法求助

第一步：训练过程中预埋指标存储逻辑 #

在训练启动前先定义4个空列表，按epoch粒度存储训练、验证阶段的损失和准确率指标，以下是适配BERT分类任务的训练循环示例：

import torch
import matplotlib.pyplot as plt

# 提前定义运行硬件
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 训练前初始化指标存储列表
train_loss_all = []
train_acc_all = []
val_loss_all = []
val_acc_all = []

num_epochs = 10 # 你的训练轮次配置

for epoch in range(num_epochs):
    model.train()
    train_loss = 0.0
    train_correct = 0
    train_total = 0

    # 训练批次循环
    for batch in train_dataloader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        logits = outputs.logits

        loss.backward()
        optimizer.step()

        # 累计训练损失
        train_loss += loss.item()
        # 计算训练准确率
        preds = torch.argmax(logits, dim=1)
        train_correct += (preds == labels).sum().item()
        train_total += labels.size(0)

    # 计算当前epoch平均训练指标，存入列表
    avg_train_loss = train_loss / len(train_dataloader)
    avg_train_acc = train_correct / train_total
    train_loss_all.append(avg_train_loss)
    train_acc_all.append(avg_train_acc)

    # 验证阶段
    model.eval()
    val_loss = 0.0
    val_correct = 0
    val_total = 0
    with torch.no_grad():
        for batch in val_dataloader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss
            logits = outputs.logits

            val_loss += loss.item()
            preds = torch.argmax(logits, dim=1)
            val_correct += (preds == labels).sum().item()
            val_total += labels.size(0)

    # 计算当前epoch平均验证指标，存入列表
    avg_val_loss = val_loss / len(val_dataloader)
    avg_val_acc = val_correct / val_total
    val_loss_all.append(avg_val_loss)
    val_acc_all.append(avg_val_acc)

    # 打印每轮指标方便核对
    print(f"Epoch {epoch+1}/{num_epochs}, 训练损失: {avg_train_loss:.4f}, 训练准确率: {avg_train_acc:.4f}, 验证损失: {avg_val_loss:.4f}, 验证准确率: {avg_val_acc:.4f}")

第二步：训练完成后绘制曲线 #

以下代码可生成和你之前LSTM任务效果一致的双列曲线图：

# 设置画布，1行2列分别展示损失、准确率
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

# 绘制损失曲线
ax1.plot(range(1, num_epochs+1), train_loss_all, label='训练损失')
ax1.plot(range(1, num_epochs+1), val_loss_all, label='验证损失')
ax1.set_title('训练/验证损失变化')
ax1.set_xlabel('Epoch')
ax1.set_ylabel('Loss')
ax1.legend()
ax1.grid(alpha=0.3)

# 绘制准确率曲线
ax2.plot(range(1, num_epochs+1), train_acc_all, label='训练准确率')
ax2.plot(range(1, num_epochs+1), val_acc_all, label='验证准确率')
ax2.set_title('训练/验证准确率变化')
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Accuracy')
ax2.legend()
ax2.grid(alpha=0.3)

plt.tight_layout()
plt.show()
# 可取消注释下方代码直接保存图片到本地
# plt.savefig('bert_train_curve.png')

注意事项 #

• 你的批次配置为batch size=3，不需要修改上述逻辑，指标计算会自动按总样本数做平均
• 如果你没有划分验证集，只需要删掉验证相关的存储和绘图代码即可
• 要确保提前安装了matplotlib依赖，没有安装的话执行pip install matplotlib

你之前LSTM任务的曲线效果参考：

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