迁移学习落地实践：医疗文本分类的预训练与微调

# TF-IDF + 逻辑回归
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

tfidf = TfidfVectorizer(max_features=5000)
X_train = tfidf.fit_transform(train_texts)
clf = LogisticRegression().fit(X_train, train_labels)  # 准确率：68.5%

# FastText 监督训练
# 准确率：72.3%

# 原始数据统计
raw_samples = 180000
raw_unique = 125000
after_dedup = 120000
after_quality_filter = 150000

# 最终用于预训练的数据
train_size = 120000
vocab_coverage = 0.85  # 医疗术语覆盖率

# 关键配置片段
from transformers import BertConfig, BertForMaskedLM, BertTokenizer

config = BertConfig.from_pretrained('hfl/chinese-roberta-wwm-ext')
config.vocab_size = tokenizer.vocab_size

model = BertForMaskedLM.from_pretrained(
    'hfl/chinese-roberta-wwm-ext',
    config=config
)

# 自定义 masking 策略
def mask_tokens(inputs, tokenizer, mlm_probability=0.15, term_mask_prob=0.20):
    """
    对医疗术语提高 masking 概率
    """
    labels = inputs.clone()
    probability_matrix = torch.full(labels.shape, mlm_probability)
    special_tokens_mask = [
        tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True)
        for val in labels.tolist()
    ]
    probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), value=0.0)

    # 识别医疗术语（这里简化处理）
    medical_terms = identify_medical_terms(inputs, tokenizer)
    for term_positions in medical_terms:
        probability_matrix[term_positions] = term_mask_prob

    masked_indices = torch.bernoulli(probability_matrix).bool()
    labels[~masked_indices] = -100
    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
    inputs[indices_replaced] = tokenizer.convert_tokens_to_ids(tokenizer.mask_token)

    return inputs, labels

# 同时预测分类和领域，共享 backbone
class MultiTaskModel(BertForSequenceClassification):
    def __init__(self, config):
        super().__init__(config)
        self.domain_classifier = torch.nn.Linear(config.hidden_size, 2)

    def forward(self, input_ids, attention_mask, labels=None, domain_labels=None):
        outputs = super().forward(input_ids, attention_mask, labels=labels)
        pooled_output = self.bert(input_ids, attention_mask=attention_mask).pooler_output
        domain_logits = self.domain_classifier(pooled_output)

        loss = None
        if labels is not None and domain_labels is not None:
            classification_loss = torch.nn.functional.cross_entropy(outputs.logits, labels)
            domain_loss = torch.nn.functional.cross_entropy(domain_logits, domain_labels)
            loss = classification_loss + 0.3 * domain_loss  # 领域损失加权

        return {'loss': loss, 'logits': outputs.logits, 'domain_logits': domain_logits}

# 数据增强示例
def augment_text(text, augment_ratio=1.5):
    """对小类数据进行增强"""
    augment_methods = [
        synonym_replace,
        back_translate,
        random_insert,
    ]
    aug_texts = [text]
    for _ in range(int(augment_ratio - 1)):
        method = random.choice(augment_methods)
        aug_texts.append(method(text))
    return aug_texts

# 类别统计和平衡
class_counts = Counter(train_labels)
min_class_size = 200
balanced_data = []

for text, label in zip(train_texts, train_labels):
    if class_counts[label] < min_class_size:
        aug_texts = augment_text(text, min_class_size / class_counts[label])
        balanced_data.extend([(t, label) for t in aug_texts])
    else:
        balanced_data.append((text, label))

from transformers import BertForSequenceClassification

model = BertForSequenceClassification.from_pretrained(
    './domain_pretrained_medical_bert',
    num_labels=10
)

training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    warmup_steps=100,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=50,
    evaluation_strategy='epoch',
    save_strategy='epoch',
    load_best_model_at_end=True,
    metric_for_best_model='f1',
    learning_rate=2e-5,
)

for param in model.bert.parameters():
    param.requires_grad = False
optimizer = AdamW(model.classifier.parameters(), lr=2e-5)
# 训练 15 分钟，显存 2.1 GB，准确率 84.3%

for param in model.parameters():
    param.requires_grad = True
optimizer = AdamW(model.parameters(), lr=2e-5)
# 训练 1.5 小时，显存 10.8 GB，准确率 88.2%

# 阶段 1：只训练分类头（3 epochs），lr=2e-5
# 阶段 2：解冻最后 2 层（3 epochs）
for param in model.bert.encoder.layer[-2:].parameters():
    param.requires_grad = True
optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-5)
# 阶段 3：全参数微调（4 epochs）
optimizer = AdamW(model.parameters(), lr=5e-6)
# 训练 1.2 小时，显存 10.8 GB，准确率 88.5%

def is_valid_sample(text):
    """数据质量过滤"""
    # 过滤广告关键词
    ad_keywords = ['广告', '推广', '优惠', '购买链接']
    if any(kw in text for kw in ad_keywords):
        return False

    # 过滤过短文本
    if len(text.strip()) < 10:
        return False

    # 过滤非中文为主的内容
    chinese_ratio = len(re.findall(r'[一-鿿]', text)) / len(text)
    if chinese_ratio < 0.7:
        return False

    return True

# 混合训练数据
domain_train = load_medical_qa_data()
general_train = load_general_qa_data()[:200]  # 少量通用数据

mixed_train = domain_train + general_train
shuffle(mixed_train)

from sklearn.metrics import precision_recall_fscore_support

def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)

    precision, recall, f1, _ = precision_recall_fscore_support(
        labels, preds, average='weighted'
    )
    macro_f1 = precision_recall_fscore_support(
        labels, preds, average='macro'
    )[2]

    return {
        'accuracy': (preds == labels).mean(),
        'f1': f1,
        'precision': precision,
        'recall': recall,
        'macro_f1': macro_f1  # 关注小类表现
    }

from transformers import get_linear_schedule_with_warmup

total_steps = len(train_loader) * num_epochs
optimizer = AdamW(model.parameters(), lr=2e-5)
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=int(0.1 * total_steps),
    num_training_steps=total_steps
)

for epoch in range(num_epochs):
    for batch in train_loader:
        optimizer.zero_grad()
        loss = model(**batch).loss
        loss.backward()
        optimizer.step()
        scheduler.step()
# 2e-5 + 10% warmup 效果最稳

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
accumulation_steps = 4  # batch size 从 32 降到 8，用梯度累积等效补回

for i, batch in enumerate(train_loader):
    optimizer.zero_grad()
    with autocast():
        loss = model(**batch).loss / accumulation_steps
    scaler.scale(loss).backward()

    if (i + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()
# 显存占用从 10.8 GB 降到 5.6 GB

class EWC:
    def __init__(self, model, dataloader):
        self.model = model
        self.fisher = self.compute_fisher(dataloader)
        self.optimal_params = {n: p.clone() for n, p in model.named_parameters()}

    def compute_fisher(self, dataloader):
        fisher = {n: torch.zeros_like(p) for n, p in self.model.named_parameters()}
        self.model.eval()
        for batch in dataloader:
            self.model(**batch).loss.backward()
            for n, p in self.model.named_parameters():
                if p.grad is not None:
                    fisher[n] += p.grad.pow(2)
        for n in fisher:
            fisher[n] /= len(dataloader)
        return fisher

    def penalty(self):
        return sum(
            (self.fisher[n] * (p - self.optimal_params[n]).pow(2)).sum()
            for n, p in self.model.named_parameters()
        )

# 训练时把 EWC 惩罚加进 loss
ewc = EWC(model, original_dataloader)
for batch in train_loader:
    loss = model(**batch).loss + 0.1 * ewc.penalty()

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["query", "value"],
    lora_dropout=0.1,
    bias="none",
    task_type="SEQ_CLS",
)
model = get_peft_model(model, lora_config)
# 可训练参数从 110M 降到 2.4M，准确率 88.1%，泛化能力更好

# 线上最终流程（简化）
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer

model = AutoModelForSequenceClassification.from_pretrained(
    './domain_pretrained_medical_bert', num_labels=10
)
model = get_peft_model(model, LoraConfig(
    r=16, lora_alpha=32,
    target_modules=["query", "key", "value"],
    lora_dropout=0.05, bias="none", task_type="SEQ_CLS",
))

training_args = TrainingArguments(
    output_dir='./final_model',
    num_train_epochs=8,
    per_device_train_batch_size=16,
    gradient_accumulation_steps=2,
    learning_rate=3e-5,
    warmup_ratio=0.1,
    weight_decay=0.01,
    fp16=True,
    evaluation_strategy='epoch',
    load_best_model_at_end=True,
    metric_for_best_model='f1',
)