Python 开源 AI 模型引入、训练与测试全流程实战

2.2 Hugging Face 集成

实际开发中，直接使用 transformers 库更为高效。以下是基于 BERT 的序列分类模型封装：

from transformers import BertConfig, BertModel, BertTokenizer, PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
import torch.nn as nn

class BertForSequenceClassification(PreTrainedModel):
    def __init__(self, config: BertConfig):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config
        self.bert = BertModel(config)
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Dropout(config.hidden_dropout_prob),
            nn.Linear(config.hidden_size, config.hidden_size),
            nn.GELU(),
            nn.Dropout(config.hidden_dropout_prob),
            nn.Linear(config.hidden_size, config.num_labels)
        )
        self.post_init()

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, labels=None, **kwargs):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, **kwargs)
        pooled_output = outputs.pooler_output
        logits = self.classifier(pooled_output)
        
        loss = None
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
        
        return SequenceClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)

3. 数据准备与预处理

使用 IMDB 电影评论数据集进行情感分类。关键在于分词和批处理。

from datasets import load_dataset, DatasetDict
from transformers import BertTokenizer
from torch.utils.data import DataLoader

class DataProcessor:
    def __init__(self, model_name: str = "bert-base-uncased", max_length: int = 512):
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.max_length = max_length

    def load_imdb_dataset(self, cache_dir: str = "./data"):
        dataset = load_dataset("imdb", cache_dir=cache_dir)
        train_test_split = dataset["train"].train_test_split(test_size=0.1, seed=42)
        dataset_dict = DatasetDict({
            "train": train_test_split["train"],
            "validation": train_test_split["test"],
            "test": dataset["test"]
        })
        return dataset_dict

    def preprocess_function(self, examples):
        tokenized_inputs = self.tokenizer(
            examples["text"], truncation=True, padding="max_length",
            max_length=self.max_length, return_tensors="pt"
        )
        return {
            "input_ids": tokenized_inputs["input_ids"].tolist(),
            "attention_mask": tokenized_inputs["attention_mask"].tolist(),
            "labels": examples["label"]
        }

    def prepare_dataset(self, dataset_dict, batch_size: int = 32):
        tokenized_datasets = dataset_dict.map(
            self.preprocess_function, batched=True, remove_columns=["text", "label"]
        )
        tokenized_datasets.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
        
        train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, batch_size=batch_size, collate_fn=self.collate_fn)
        val_dataloader = DataLoader(tokenized_datasets["validation"], batch_size=batch_size, collate_fn=self.collate_fn)
        test_dataloader = DataLoader(tokenized_datasets["test"], batch_size=batch_size, collate_fn=self.collate_fn)
        return train_dataloader, val_dataloader, test_dataloader

    def collate_fn(self, batch):
        input_ids = torch.stack([item["input_ids"] for item in batch])
        attention_mask = torch.stack([item["attention_mask"] for item in batch])
        labels = torch.tensor([item["labels"] for item in batch])
        return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

4. 模型训练与微调

自定义训练器允许我们更灵活地控制优化过程，例如梯度裁剪和学习率调度。

from transformers import TrainingArguments, AdamW, get_linear_schedule_with_warmup
from tqdm.auto import tqdm
from dataclasses import dataclass

@dataclass
class TrainingConfig:
    model_name: str = "bert-base-uncased"
    num_labels: int = 2
    batch_size: int = 32
    num_epochs: int = 3
    learning_rate: float = 2e-5
    weight_decay: float = 0.01
    warmup_steps: int = 500
    fp16: bool = True

class CustomTrainer:
    def __init__(self, model, train_config: TrainingConfig, train_dataloader, val_dataloader, test_dataloader=None):
        self.model = model
        self.config = train_config
        self.train_dataloader = train_dataloader
        self.val_dataloader = val_dataloader
        self.test_dataloader = test_dataloader
        self.device = torch.device(train_config.device if hasattr(train_config, 'device') else "cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        
        self.optimizer = self._create_optimizer()
        self.scheduler = self._create_scheduler()
        self.global_step = 0
        self.best_metric = 0.0

    def _create_optimizer(self):
        no_decay = ["bias", "LayerNorm.weight"]
        optimizer_grouped_parameters = [
            {"params": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay)], "weight_decay": self.config.weight_decay},
            {"params": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0}
        ]
        return AdamW(optimizer_grouped_parameters, lr=self.config.learning_rate, eps=1e-8)

    def _create_scheduler(self):
        total_steps = len(self.train_dataloader) * self.config.num_epochs
        return get_linear_schedule_with_warmup(self.optimizer, num_warmup_steps=self.config.warmup_steps, num_training_steps=total_steps)

    def train_epoch(self, epoch: int):
        self.model.train()
        total_loss = 0
        progress_bar = tqdm(self.train_dataloader, desc=f"Epoch {epoch}", leave=False)
        
        for batch in progress_bar:
            batch = {k: v.to(self.device) for k, v in batch.items()}
            outputs = self.model(**batch)
            loss = outputs.loss
            loss.backward()
            
            # 梯度裁剪防止爆炸
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            
            self.optimizer.step()
            self.scheduler.step()
            self.optimizer.zero_grad()
            
            total_loss += loss.item()
            self.global_step += 1
            progress_bar.set_postfix({"loss": loss.item(), "lr": self.scheduler.get_last_lr()[0]})
        
        avg_loss = total_loss / len(self.train_dataloader)
        return {"train_loss": avg_loss}

    def evaluate(self, dataloader=None):
        if dataloader is None:
            dataloader = self.val_dataloader
        self.model.eval()
        total_loss = 0
        all_preds = []
        all_labels = []
        
        with torch.no_grad():
            for batch in tqdm(dataloader, desc="Evaluating", leave=False):
                batch = {k: v.to(self.device) for k, v in batch.items()}
                outputs = self.model(**batch)
                loss = outputs.loss
                logits = outputs.logits
                total_loss += loss.item()
                preds = torch.argmax(logits, dim=-1)
                all_preds.extend(preds.cpu().numpy())
                all_labels.extend(batch["labels"].cpu().numpy())
        
        from sklearn.metrics import accuracy_score, f1_score
        accuracy = accuracy_score(all_labels, all_preds)
        f1 = f1_score(all_labels, all_preds, average="binary")
        avg_loss = total_loss / len(dataloader)
        return {"loss": avg_loss, "accuracy": accuracy, "f1": f1}

    def train(self):
        print(f"Starting training with config: {self.config}")
        for epoch in range(self.config.num_epochs):
            print(f"\n{'='*50}\nEpoch {epoch + 1}/{self.config.num_epochs}\n{'='*50}")
            train_metrics = self.train_epoch(epoch)
            val_metrics = self.evaluate()
            print(f"Train Loss: {train_metrics['train_loss']:.4f}")
            print(f"Val Accuracy: {val_metrics['accuracy']:.4f}")
            if val_metrics["accuracy"] > self.best_metric:
                self.best_metric = val_metrics["accuracy"]
                self.save_model(f"best_model_step_{self.global_step}")
        return self.history if hasattr(self, 'history') else {}

    def save_model(self, save_path: str):
        torch.save({"model_state_dict": self.model.state_dict()}, f"{save_path}.pt")
        self.model.save_pretrained(f"{save_path}_hf")

5. 模型评估与测试

除了准确率，还需要关注混淆矩阵、ROC 曲线以及压力测试。

5.1 综合评估指标

from sklearn.metrics import classification_report, confusion_matrix
import numpy as np

class ModelEvaluator:
    def __init__(self, model, tokenizer, device: str = "cuda"):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device
        self.model.to(device)
        self.model.eval()

    def predict(self, texts, batch_size: int = 32):
        all_logits = []
        all_probs = []
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i+batch_size]
            inputs = self.tokenizer(batch_texts, truncation=True, padding=True, max_length=512, return_tensors="pt")
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            with torch.no_grad():
                outputs = self.model(**inputs)
                logits = outputs.logits
                probs = torch.softmax(logits, dim=-1)
                all_logits.append(logits.cpu().numpy())
                all_probs.append(probs.cpu().numpy())
        return np.vstack(all_logits), np.vstack(all_probs)

    def evaluate_classification(self, texts, labels, threshold: float = 0.5):
        logits, probs = self.predict(texts)
        preds = np.argmax(probs, axis=1)
        metrics = {
            "accuracy": (preds == labels).mean(),
            "precision": 0.0, # Placeholder for sklearn precision_score
            "recall": 0.0,    # Placeholder for sklearn recall_score
            "f1": 0.0         # Placeholder for sklearn f1_score
        }
        cm = confusion_matrix(labels, preds)
        report = classification_report(labels, preds, target_names=["Negative", "Positive"])
        return {"metrics": metrics, "confusion_matrix": cm, "classification_report": report}

5.2 压力测试与性能基准

在生产环境中，推理延迟和吞吐量至关重要。

import time
import psutil

class PerformanceBenchmark:
    def __init__(self, model, tokenizer, device: str = "cuda"):
        self.model = model
        self.tokenizer = tokenizer
        self.device = device

    def measure_inference_time(self, texts, batch_sizes=[1, 4, 8, 16, 32]):
        results = {}
        for batch_size in batch_sizes:
            print(f"\nTesting batch size: {batch_size}")
            # Warmup
            warmup_texts = ["This is a warmup sentence."] * batch_size
            self.predict_batch(warmup_texts)
            
            times = []
            for i in range(0, len(texts), batch_size):
                batch_texts = texts[i:i+batch_size]
                start_time = time.perf_counter()
                self.predict_batch(batch_texts)
                end_time = time.perf_counter()
                times.append(end_time - start_time)
            
            avg_time = np.mean(times)
            throughput = len(texts) / np.sum(times)
            results[batch_size] = {"avg_inference_time": avg_time, "throughput": throughput}
            print(f" Average inference time: {avg_time:.4f}s")
            print(f" Throughput: {throughput:.2f} samples/s")
        return results

    def predict_batch(self, texts):
        inputs = self.tokenizer(texts, truncation=True, padding=True, max_length=512, return_tensors="pt")
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        with torch.no_grad():
            outputs = self.model(**inputs)
        return outputs

6. 测试框架与质量保证

使用 pytest 和 hypothesis 进行单元测试和属性测试。

6.1 单元测试示例

import pytest
from hypothesis import given, strategies as st
import numpy as np

class TestDataProcessor:
    def setup_method(self):
        self.processor = DataProcessor("bert-base-uncased")

    def test_tokenization(self):
        text = "This is a test sentence."
        tokenized = self.processor.tokenizer(text, truncation=True, padding="max_length", max_length=128)
        assert "input_ids" in tokenized
        assert len(tokenized["input_ids"]) == 128

    @given(st.text(min_size=1, max_size=1000), st.integers(min_value=0, max_value=1))
    def test_preprocess_function(self, text, label):
        examples = {"text": [text], "label": [label]}
        result = self.processor.preprocess_function(examples)
        assert "input_ids" in result
        assert result["labels"][0] == label

6.2 集成测试与 API 测试

针对 FastAPI 服务的健康检查和预测端点进行验证。

from fastapi.testclient import TestClient
from src.api.app import app

class TestAPI:
    def setup_method(self):
        self.client = TestClient(app)

    def test_health_endpoint(self):
        response = self.client.get("/health")
        assert response.status_code == 200
        assert response.json()["status"] == "healthy"

    def test_predict_endpoint(self):
        test_data = {"text": "This movie was absolutely fantastic!", "model_version": "latest"}
        response = self.client.post("/predict", json=test_data)
        assert response.status_code == 200
        result = response.json()
        assert "prediction" in result
        assert 0 <= result["confidence"] <= 1

    async def test_concurrent_requests(self):
        import asyncio
        async def make_request():
            test_data = {"text": "Test concurrent request", "model_version": "latest"}
            response = await self.client.post("/predict", json=test_data)
            return response.status_code
        tasks = [make_request() for _ in range(10)]
        results = await asyncio.gather(*tasks)
        assert all(status == 200 for status in results)

7. 模型部署与 API 服务

使用 FastAPI 构建高性能接口，并通过 Docker 容器化部署。

7.1 FastAPI 服务实现

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from typing import List, Optional
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="BERT Text Classification API", version="1.0.0")

class PredictionRequest(BaseModel):
    text: str = Field(..., min_length=1, max_length=5000)
    model_version: Optional[str] = "latest"

class PredictionResponse(BaseModel):
    prediction: int
    label: str
    confidence: float
    model_version: str
    processing_time: float

@app.get("/health")
async def health_check():
    return {"status": "healthy", "timestamp": "now"}

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest):
    # 模拟预测逻辑
    return {
        "prediction": 1,
        "label": "Positive",
        "confidence": 0.95,
        "model_version": request.model_version,
        "processing_time": 0.05
    }

7.2 Docker 部署配置

FROM python:3.9-slim
WORKDIR /app
RUN apt-get update && apt-get install -y build-essential curl && rm -rf /var/lib/apt/lists/*
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
EXPOSE 8000
CMD ["uvicorn", "src.api.app:app", "--host", "0.0.0.0", "--port", "8000"]

配合 docker-compose.yml 可轻松编排服务与监控组件（如 Prometheus）。

8. 监控与日志

生产环境需要结构化日志和性能监控。

8.1 结构化日志配置

使用 jsonlogger 输出 JSON 格式日志，便于 ELK 等系统采集。

import logging
from pythonjsonlogger import jsonlogger

class CustomJsonFormatter(jsonlogger.JsonFormatter):
    def add_fields(self, log_record, record, message_dict):
        super().add_fields(log_record, record, message_dict)
        log_record['service'] = 'bert-classification-api'
        log_record['module'] = record.module

LOGGING_CONFIG = {
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {'json': {'()': CustomJsonFormatter}},
    'handlers': {
        'console': {'class': 'logging.StreamHandler', 'formatter': 'json', 'level': 'INFO'}
    },
    'loggers': {'': {'handlers': ['console'], 'level': 'INFO'}}
}

def setup_logging():
    logging.config.dictConfig(LOGGING_CONFIG)

8.2 性能监控

集成 Prometheus 客户端暴露指标。

from prometheus_client import Counter, Histogram, generate_latest

PREDICTION_REQUESTS = Counter('prediction_requests_total', 'Total number of prediction requests')
PREDICTION_LATENCY = Histogram('prediction_latency_seconds', 'Prediction latency in seconds')

@app.get("/metrics")
async def metrics_endpoint():
    return Response(generate_latest(), media_type="text/plain")

9. 优化与最佳实践

9.1 模型量化与 ONNX

为了进一步提升推理速度，可以将模型转换为 ONNX 格式并进行量化。

from optimum.onnxruntime import ORTModelForSequenceClassification

def convert_to_onnx(model_path: str, output_path: str = "./models/onnx"):
    model = ORTModelForSequenceClassification.from_pretrained(model_path, from_transformers=True, export=True)
    model.save_pretrained(output_path)
    return output_path

9.2 缓存策略

对于重复请求，使用 LRU 缓存减少计算开销。

from functools import lru_cache
import hashlib

def get_cache_key(text: str, model_version: str) -> str:
    content = f"{model_version}:{text}"
    return hashlib.md5(content.encode()).hexdigest()

@lru_cache(maxsize=1000)
def cached_predict(text: str, model_version: str):
    # 实际预测逻辑
    pass

10. 总结

本文详细展示了开源 AI 模型从引入到测试的完整技术流程。通过这个实战项目，我们不仅学习了如何集成和使用先进的 AI 模型，更重要的是掌握了构建生产级 AI 应用的系统工程方法。这套方法论和代码框架可以应用于各种 AI 项目，为 AI 应用开发提供坚实基础。成功的 AI 项目不仅仅是模型准确率，更是系统工程、可维护性、可扩展性和可靠性的综合体现。