压缩大型语言模型(LLMs)：缩小10倍、性能保持不变

from datasets import load_dataset  
 
 from transformers import AutoTokenizer, AutoModelForSequenceClassification  
 from transformers import DistilBertForSequenceClassification, DistilBertConfig  
 
 import torch  
 import torch.nn as nn  
 import torch.optim as optim  
 from torch.utils.data import DataLoader  
 
 from sklearn.metrics import accuracy_score, precision_recall_fscore_support

data = load_dataset("shawhin/phishing-site-classification")

# 使用Nvidia GPU  
 device = torch.device('cuda')  
 
 # 加载教师模型和分词器
 model_path = "shawhin/bert-phishing-classifier_teacher"  
 
 tokenizer = AutoTokenizer.from_pretrained(model_path)  
 teacher_model = AutoModelForSequenceClassification.from_pretrained(model_path)  
                                                  .to(device)

# 加载学生模型
 my_config = DistilBertConfig(n_heads=8, n_layers=4) # 每层减少4个头,总共减少2层  
 
 student_model = DistilBertForSequenceClassification  
                                    .from_pretrained("distilbert-base-uncased",  
                                     config=my_config,)  
                                    .to(device)

# 定义文本预处理函数
 def preprocess_function(examples):  
     return tokenizer(examples["text"], padding='max_length', truncation=True)  
 
 # 对所有数据集进行标记化
 tokenized_data = data.map(preprocess_function, batched=True)  
 tokenized_data.set_format(type='torch',  
                           columns=['input_ids', 'attention_mask', 'labels'])

# 评估模型性能的函数
 def evaluate_model(model, dataloader, device):  
     model.eval()  # 将模型设置为评估模式
     all_preds = []  
     all_labels = []  
 
     # 禁用梯度计算
     with torch.no_grad():  
         for batch in dataloader:  
             input_ids = batch['input_ids'].to(device)  
             attention_mask = batch['attention_mask'].to(device)  
             labels = batch['labels'].to(device)  
 
             # 前向传播获取logits  
             outputs = model(input_ids, attention_mask=attention_mask)  
             logits = outputs.logits  
 
             # 获取预测结果
             preds = torch.argmax(logits, dim=1).cpu().numpy()  
             all_preds.extend(preds)  
             all_labels.extend(labels.cpu().numpy())  
 
     # 计算评估指标
     accuracy = accuracy_score(all_labels, all_preds)  
     precision, recall, f1, _ = precision_recall_fscore_support(all_labels,  
                                                               all_preds,  
                                                 average='binary')  
 
     return accuracy, precision, recall, f1

# 计算蒸馏损失和硬标签损失的函数
 def distillation_loss(student_logits, teacher_logits,  
                       true_labels, temperature, alpha):  
     # 从教师logits计算软目标
     soft_targets = nn.functional.softmax(teacher_logits / temperature, dim=1)  
     student_soft = nn.functional.log_softmax(student_logits / temperature, dim=1)  
 
     # 蒸馏的KL散度损失
     distill_loss = nn.functional.kl_div(student_soft,  
                                     soft_targets,  
                                     reduction='batchmean') * (temperature ** 2)  
 
     # 硬标签的交叉熵损失
     hard_loss = nn.CrossEntropyLoss()(student_logits, true_labels)  
 
     # 结合损失
     loss = alpha * distill_loss + (1.0 - alpha) * hard_loss  
 
     return loss

# 超参数
 batch_size = 32  
 lr = 1e-4  
 num_epochs = 5  
 temperature = 2.0  
 alpha = 0.5  
 
 # 定义优化器
 optimizer = optim.Adam(student_model.parameters(), lr=lr)  
 
 # 创建训练数据加载器
 dataloader = DataLoader(tokenized_data['train'], batch_size=batch_size)  
 # 创建测试数据加载器
 test_dataloader = DataLoader(tokenized_data['test'], batch_size=batch_size)

# 将学生模型设置为训练模式
 student_model.train()  
 
 # 训练模型
 for epoch in range(num_epochs):  
     for batch in dataloader:  
         # 准备输入
         input_ids = batch['input_ids'].to(device)  
         attention_mask = batch['attention_mask'].to(device)  
         labels = batch['labels'].to(device)  
 
         # 禁用教师模型的梯度计算
         with torch.no_grad():  
             teacher_outputs = teacher_model(input_ids,  
                                             attention_mask=attention_mask)  
             teacher_logits = teacher_outputs.logits  
 
         # 学生模型前向传播
         student_outputs = student_model(input_ids,  
                                         attention_mask=attention_mask)  
         student_logits = student_outputs.logits  
 
         # 计算蒸馏损失
         loss = distillation_loss(student_logits, teacher_logits, labels,  
                                   temperature, alpha)  
 
         # 反向传播
         optimizer.zero_grad()  
         loss.backward()  
         optimizer.step()  
 
     print(f"第 {epoch + 1} 轮训练完成,损失: {loss.item()}")  
 
     # 评估教师模型
     teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 =  
                          evaluate_model(teacher_model, test_dataloader, device)  
 
     print(f"教师模型 (测试集) - 准确率: {teacher_accuracy:.4f},  
                                 精确率: {teacher_precision:.4f},  
                                 召回率: {teacher_recall:.4f},  
                                 F1分数: {teacher_f1:.4f}")  
 
     # 评估学生模型
     student_accuracy, student_precision, student_recall, student_f1 =  
                          evaluate_model(student_model, test_dataloader, device)  
       
     print(f"学生模型 (测试集) - 准确率: {student_accuracy:.4f},  
                                 精确率: {student_precision:.4f},  
                                 召回率: {student_recall:.4f},  
                                 F1分数: {student_f1:.4f}")  
     print("\n")  
 
     # 将学生模型重新设置为训练模式
     student_model.train()

# 创建验证数据加载器
 validation_dataloader = DataLoader(tokenized_data['validation'], batch_size=8)  
 
 # 评估教师模型
 teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 =  
                    evaluate_model(teacher_model, validation_dataloader, device)  
 print(f"教师模型 (验证集) - 准确率: {teacher_accuracy:.4f},  
                            精确率: {teacher_precision:.4f},  
                            召回率: {teacher_recall:.4f},  
                            F1分数: {teacher_f1:.4f}")  
 
 # 评估学生模型
 student_accuracy, student_precision, student_recall, student_f1 =  
                    evaluate_model(student_model, validation_dataloader, device)  
 print(f"学生模型 (验证集) - 准确率: {student_accuracy:.4f},  
                            精确率: {student_precision:.4f},  
                            召回率: {student_recall:.4f},  
                            F1分数: {student_f1:.4f}")

from transformers import BitsAndBytesConfig  
 
 # 以4位精度加载模型
 nf4_config = BitsAndBytesConfig(  
     load_in_4bit=True,  
     bnb_4bit_quant_type="nf4",  
     bnb_4bit_compute_dtype = torch.bfloat16,  
     bnb_4bit_use_double_quant=True  
 )  
 
 model_nf4 = AutoModelForSequenceClassification.from_pretrained(model_id,  
                                                 device_map=device,  
                                                 quantization_config=nf4_config)

# 评估量化后的学生模型
 quantized_accuracy, quantized_precision, quantized_recall, quantized_f1 =  
                        evaluate_model(model_nf4, validation_dataloader, device)  
 
 print("量化后性能")  
 print(f"准确率: {quantized_accuracy:.4f},  
         精确率: {quantized_precision:.4f},  
         召回率: {quantized_recall:.4f},  
         F1分数: {quantized_f1:.4f}")