Transformer 革命：知識蒸餾訓練器

from transformers import TrainingArguments

class DistillationTrainingArguments(TrainingArguments):
    def __init__(self, *args, alpha=0.5, temperature=2.0, **kwargs):
        super().__init__(*args, **kwargs)
        self.alpha = alpha
        self.temperature = temperature

import torch.nn as nn
import torch.nn.functional as F
from transformers import Trainer

class DistillationTrainer(Trainer):
    def __init__(self, *args, teacher_model=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.teacher_model = teacher_model
        
    def compute_loss(self, model, inputs, return_outputs=False):
        outputs_stu = model(**inputs)
        # 提取學生模型的交叉熵損失和logits
        loss_ce = outputs_stu.loss
        logits_stu = outputs_stu.logits
        
        # 提取教師模型的logits
        with torch.no_grad():
            outputs_tea = self.teacher_model(**inputs)
            logits_tea = outputs_tea.logits
            
        # 軟化機率分佈並計算蒸餾損失
        loss_fct = nn.KLDivLoss(reduction="batchmean")
        loss_kd = self.args.temperature ** 2 * loss_fct(
            F.log_softmax(logits_stu / self.args.temperature, dim=-1),
            F.softmax(logits_tea / self.args.temperature, dim=-1))
            
        # 回傳加權後的學生損失
        loss = self.args.alpha * loss_ce + (1. - self.args.alpha) * loss_kd
        return (loss, outputs_stu) if return_outputs else loss

from transformers import AutoTokenizer

student_ckpt = "distilbert-base-uncased"
student_tokenizer = AutoTokenizer.from_pretrained(student_ckpt)

def tokenize_text(batch):
    return student_tokenizer(batch["text"], truncation=True)

clinc_enc = clinc.map(tokenize_text, batched=True, remove_columns=["text"])
clinc_enc = clinc_enc.rename_column("intent", "labels")

def compute_metrics(pred):
    predictions, labels = pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy_score.compute(predictions=predictions, references=labels)

batch_size = 48
finetuned_ckpt = "distilbert-base-uncased-finetuned-clinc"
student_training_args = DistillationTrainingArguments(
    output_dir=finetuned_ckpt, 
    evaluation_strategy="epoch",
    num_train_epochs=5, 
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size, 
    alpha=1, 
    weight_decay=0.01,
    push_to_hub=True)

# 提取教師和學生的隱藏狀態
hidden_states_stu = outputs_stu.hidden_states[-1]  # 最後一層的隱藏狀態
hidden_states_tea = outputs_tea.hidden_states[-1]

# 計算餘弦相似度損失
cosine_loss = 1 - F.cosine_similarity(hidden_states_stu, hidden_states_tea, dim=-1).mean()

# 加入到總損失
loss = self.args.alpha * loss_ce + self.args.beta * loss_kd + self.args.gamma * cosine_loss

id2label = pipe.model.config.id2label
label2id = pipe.model.config.label2id

from transformers import AutoConfig

num_labels = intents.num_classes
student_config = (AutoConfig
    .from_pretrained(student_ckpt, num_labels=num_labels,
                    id2label=id2label, label2id=label2id))

import torch
from transformers import AutoModelForSequenceClassification

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def student_init():
    return (AutoModelForSequenceClassification
            .from_pretrained(student_ckpt, config=student_config).to(device))

teacher_ckpt = "transformersbook/bert-base-uncased-finetuned-clinc"
teacher_model = (AutoModelForSequenceClassification
                .from_pretrained(teacher_ckpt, num_labels=num_labels)
                .to(device))

distilbert_trainer = DistillationTrainer(model_init=student_init,
                                       teacher_model=teacher_model, 
                                       args=student_training_args,
                                       train_dataset=clinc_enc['train'], 
                                       eval_dataset=clinc_enc['validation'],
                                       compute_metrics=compute_metrics, 
                                       tokenizer=student_tokenizer)

distilbert_trainer.train()

distilbert_trainer.push_to_hub("Training completed!")

finetuned_ckpt = "transformersbook/distilbert-base-uncased-finetuned-clinc"
pipe = pipeline("text-classification", model=finetuned_ckpt)

optim_type = "DistilBERT"
pb = PerformanceBenchmark(pipe, clinc["test"], optim_type=optim_type)
perf_metrics.update(pb.run_benchmark())

import pandas as pd

def plot_metrics(perf_metrics, current_optim_type):
    df = pd.DataFrame.from_dict(perf_metrics, orient='index')
    for idx in df.index:
        df_opt = df.loc[idx]
        # 為當前最佳化型別增加虛線圓圈
        if idx == current_optim_type:
            plt.scatter(df_opt["time_avg_ms"], df_opt["accuracy"] * 100,
                       alpha=0.5, s=df_opt["size_mb"], label=idx,
                       marker='$\u25CC$')
        else:
            plt.scatter(df_opt["time_avg_ms"], df_opt["accuracy"] * 100,
                       s=df_opt["size_mb"], label=idx, alpha=0.5)
    legend = plt.legend(bbox_to_anchor=(1,1))
    for handle in legend.legendHandles:
        handle.set_sizes([20])
    plt.ylim(80,90)
    # 使用最慢的模型來定義X軸範圍
    xlim = int(perf_metrics["BERT baseline"]["time_avg_ms"] + 3)
    plt.xlim(1, xlim)
    plt.ylabel("Accuracy (%)")
    plt.xlabel("Average latency (ms)")
    plt.show()

plot_metrics(perf_metrics, optim_type)

f(x, y) = (1 - x)² + 100(y - x²)²

def objective(trial):
    x = trial.suggest_float("x", -2, 2)
    y = trial.suggest_float("y", -2, 2)
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2

import optuna
study = optuna.create_study()
study.optimize(objective, n_trials=1000)

print(study.best_params)
# 輸出: {'x': 1.003024865971437, 'y': 1.00315167589307}

def hp_space(trial):
    return {"num_train_epochs": trial.suggest_int("num_train_epochs", 5, 10),
            "alpha": trial.suggest_float("alpha", 0, 1),
            "temperature": trial.suggest_int("temperature", 2, 20)}

best_run = distilbert_trainer.hyperparameter_search(
    n_trials=20, direction="maximize", hp_space=hp_space)

print(best_run)
# 輸出: BestRun(run_id='1', objective=0.927741935483871,
#        hyperparameters={'num_train_epochs': 10, 'alpha': 0.12468168730193585,
#                        'temperature': 7})

for k,v in best_run.hyperparameters.items():
    setattr(student_training_args, k, v)

# 定義一個新的倉函式庫來儲存蒸餾模型
distilled_ckpt = "distilbert-base-uncased-distilled-clinc"
student_training_args.output_dir = distilled_ckpt

# 建立一個具有最佳引數的新Trainer
distil_trainer = DistillationTrainer(model_init=student_init,
                                   teacher_model=teacher_model, 
                                   args=student_training_args,
                                   train_dataset=clinc_enc['train'], 
                                   eval_dataset=clinc_enc['validation'],
                                   compute_metrics=compute_metrics, 
                                   tokenizer=student_tokenizer)

distil_trainer.train()

distil_trainer.push_to_hub("Training complete")

distilled_ckpt = "transformersbook/distilbert-base-uncased-distilled-clinc"
pipe = pipeline("text-classification", model=distilled_ckpt)

optim_type = "Distillation"
pb = PerformanceBenchmark(pipe, clinc["test"], optim_type=optim_type)
perf_metrics.update(pb.run_benchmark())

plot_metrics(perf_metrics, optim_type)