model-training

模型训练

此技能可让 AI Agent 在结构化或非结构化数据集上训练机器学习模型。它涵盖完整的训练生命周期：加载和预处理数据、定义模型架构、配置优化器和损失函数、运行带验证的训练循环、应用学习率调度，以及保存检查点。Agent 可使用 PyTorch、TensorFlow 和 scikit-learn 等框架处理经典机器学习和深度学习工作流。

工作流

1. 加载并检查数据： 从磁盘、数据库或远程存储读取数据集。分析数据，了解特征分布、类别平衡、缺失值和数据类型。如果存在类别不平衡，则使用分层抽样将数据划分为训练集、验证集和测试集。

2. 预处理和转换： 应用特征工程，例如归一化、标准化、分词（文本）或数据增强（图像）。构建可复现且可序列化的预处理流水线，确保推理时应用相同的转换。

3. 定义模型架构： 选择或构建适合任务的模型架构。对于经典机器学习，选择梯度提升或 SVM 等估计器。对于深度学习，定义层、激活函数和正则化方式，例如 dropout 或权重衰减。适用迁移学习时，加载预训练主干网络并添加任务专用的输出头。

4. 配置训练： 设置优化器（Adam、SGD、AdamW）、损失函数（交叉熵、MSE、焦点损失）、学习率调度方案（余弦退火、阶梯衰减、预热）和批量大小。在 GPU 上训练时，可启用 torch.amp 或 tf.keras.mixed_precision 进行混合精度训练，以减少内存占用并加快计算。

5. 运行带验证的训练循环： 按指定轮数训练，并按批次或轮次记录训练损失和指标。定期在验证集上评估。实现提前停止，以便在验证性能连续达到平台期且持续轮数达到可配置值（耐心值）时停止训练。

6. 保存检查点并导出： 在验证分数最佳时以及定期保存模型检查点。将最终模型导出为便于移植的格式（ONNX、TorchScript、SavedModel），以供后续部署。将所有超参数和指标记录到 MLflow 或 Weights & Biases 等实验跟踪工具中。

支持的技术

• 框架： PyTorch、TensorFlow/Keras、scikit-learn、XGBoost、LightGBM
• 分布式训练： PyTorch DDP、Horovod、TensorFlow MirroredStrategy
• 实验跟踪： MLflow、Weights & Biases、TensorBoard
• 混合精度： torch.amp、tf.keras.mixed_precision
• 数据加载： PyTorch DataLoader、tf.data、pandas、Hugging Face Datasets

用法

向 Agent 提供数据集位置、目标变量或任务描述，以及任何约束（框架偏好、计算预算、目标指标）。Agent 将执行完整的训练工作流，并返回训练好的模型工件和评估指标。

示例

示例 1：使用 PyTorch 训练文本分类器

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from collections import Counter

# Simulated tokenized text data: 2000 samples, sequence length 50, vocab size 5000
X = torch.randint(0, 5000, (2000, 50))
y_raw = ["positive"] * 1000 + ["negative"] * 1000
le = LabelEncoder()
y = torch.tensor(le.fit_transform(y_raw), dtype=torch.long)

X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, stratify=y, random_state=42)
train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=64, shuffle=True)
val_loader = DataLoader(TensorDataset(X_val, y_val), batch_size=64)

class TextClassifier(nn.Module):
    def __init__(self, vocab_size=5000, embed_dim=128, num_classes=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, 64, batch_first=True, bidirectional=True)
        self.dropout = nn.Dropout(0.3)
        self.fc = nn.Linear(128, num_classes)

    def forward(self, x):
        x = self.embedding(x)
        _, (hidden, _) = self.lstm(x)
        hidden = torch.cat((hidden[-2], hidden[-1]), dim=1)
        return self.fc(self.dropout(hidden))

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TextClassifier().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
criterion = nn.CrossEntropyLoss()

best_val_acc, patience, patience_counter = 0.0, 3, 0
for epoch in range(10):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        loss = criterion(model(xb), yb)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    scheduler.step()

    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for xb, yb in val_loader:
            xb, yb = xb.to(device), yb.to(device)
            correct += (model(xb).argmax(1) == yb).sum().item()
            total += yb.size(0)
    val_acc = correct / total
    print(f"Epoch {epoch+1}: val_acc={val_acc:.4f}")

    if val_acc > best_val_acc:
        best_val_acc = val_acc
        torch.save(model.state_dict(), "best_model.pt")
        patience_counter = 0
    else:
        patience_counter += 1
        if patience_counter >= patience:
            print("Early stopping triggered.")
            break

示例 2：使用 Hugging Face Transformers 微调预训练模型

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
import numpy as np
from sklearn.metrics import accuracy_score, f1_score

dataset = load_dataset("imdb")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize(batch):
    return tokenizer(batch["text"], padding="max_length", truncation=True, max_length=256)

tokenized = dataset.map(tokenize, batched=True)
tokenized.set_format("torch", columns=["input_ids", "attention_mask", "label"])

model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)

def compute_metrics(eval_pred):
    preds = np.argmax(eval_pred.predictions, axis=1)
    return {"accuracy": accuracy_score(eval_pred.label_ids, preds), "f1": f1_score(eval_pred.label_ids, preds)}

training_args = TrainingArguments(
    output_dir="./results", num_train_epochs=3, per_device_train_batch_size=16,
    per_device_eval_batch_size=32, eval_strategy="epoch", save_strategy="epoch",
    load_best_model_at_end=True, metric_for_best_model="f1", fp16=True,
    learning_rate=2e-5, weight_decay=0.01, warmup_steps=500, logging_steps=100,
)

trainer = Trainer(model=model, args=training_args, train_dataset=tokenized["train"],
                  eval_dataset=tokenized["test"], compute_metrics=compute_metrics)
trainer.train()
trainer.save_model("./best_model")

最佳实践

• 处理类别不平衡的数据集时，始终使用分层划分，确保每个划分都能反映真实的类别分布。
• 微调预训练模型时使用学习率预热，避免在训练初期发生灾难性遗忘。
• 在 GPU 训练中启用混合精度（fp16 或 bf16），可将内存用量大致减半并提升吞吐量。
• 将所有信息记录到实验跟踪器中，包括超参数、每个周期的指标、系统资源用量，以及训练代码的 git 哈希值。
• 频繁保存检查点，并始终保留验证集表现最佳的检查点，避免因崩溃或训练过度而丢失进度。
• 在留出集上进行验证。该数据集从未用于训练或超参数选择，可帮助获得无偏的泛化能力估计。

边缘情况

• 小型数据集（< 1000 个样本）： 使用 k 折交叉验证，而非单次训练集／验证集划分。优先选择迁移学习或预训练模型，而不是从头开始训练。
• 类别极度不平衡（比例 > 100:1）： 使用过采样（SMOTE）、类别加权损失函数或焦点损失。评估时应采用 F1、精确率-召回率 AUC 或 Matthews 相关系数，而不是准确率。
• 训练发散或损失为 NaN： 降低学习率，应用梯度裁剪（torch.nn.utils.clip_grad_norm_），检查数据问题（如无穷值），或禁用混合精度以排查数值不稳定问题。
• 内存不足错误： 减小批次大小、启用梯度累积、使用混合精度，或改用梯度检查点，以计算量换取内存。
• 非平稳数据（概念漂移）： 定期使用新数据重新训练，按时间划分训练集和验证集，而非随机划分，并监控生产环境指标是否出现退化。