0
Jump into concepts and practice - no test required
Recommended
import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from transformers import BertTokenizer, BertForSequenceClassification, DistilBertForSequenceClassification # Assume we have a dataset class and data loaders # teacher_model: pretrained large BERT # student_model: smaller BERT or DistilBERT class DistillationLoss(nn.Module): def __init__(self, temperature=2.0, alpha=0.5): super().__init__() self.temperature = temperature self.alpha = alpha self.kl_div = nn.KLDivLoss(reduction='batchmean') self.ce = nn.CrossEntropyLoss() def forward(self, student_logits, teacher_logits, labels): soft_targets = nn.functional.log_softmax(student_logits / self.temperature, dim=1) soft_labels = nn.functional.softmax(teacher_logits / self.temperature, dim=1) distill_loss = self.kl_div(soft_targets, soft_labels) * (self.temperature ** 2) student_loss = self.ce(student_logits, labels) return self.alpha * student_loss + (1 - self.alpha) * distill_loss # Load teacher model (large) teacher_model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2) teacher_model.eval() # Load student model (smaller) student_model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=2) optimizer = optim.Adam(student_model.parameters(), lr=5e-5) criterion = DistillationLoss(temperature=2.0, alpha=0.5) # Training loop for distillation for epoch in range(3): student_model.train() for batch in train_loader: inputs = batch['input_ids'] attention_mask = batch['attention_mask'] labels = batch['labels'] with torch.no_grad(): teacher_logits = teacher_model(inputs, attention_mask=attention_mask).logits student_logits = student_model(inputs, attention_mask=attention_mask).logits loss = criterion(student_logits, teacher_logits, labels) optimizer.zero_grad() loss.backward() optimizer.step() # Evaluate student model accuracy on validation set student_model.eval() correct = 0 total = 0 with torch.no_grad(): for batch in val_loader: inputs = batch['input_ids'] attention_mask = batch['attention_mask'] labels = batch['labels'] outputs = student_model(inputs, attention_mask=attention_mask).logits preds = torch.argmax(outputs, dim=1) correct += (preds == labels).sum().item() total += labels.size(0) val_accuracy = correct / total * 100 # Apply post-training dynamic quantization quantized_model = torch.quantization.quantize_dynamic( student_model, {nn.Linear}, dtype=torch.qint8 ) # Measure model size and inference time import time import os def model_size(model): torch.save(model.state_dict(), 'temp.pth') size = os.path.getsize('temp.pth') / 1e6 # MB os.remove('temp.pth') return size size_before = model_size(student_model) size_after = model_size(quantized_model) # Inference time measurement sample_batch = next(iter(val_loader)) input_sample = sample_batch['input_ids'] attention_sample = sample_batch['attention_mask'] start = time.time() _ = student_model(input_sample, attention_mask=attention_sample) end = time.time() inference_time_before = (end - start) / input_sample.size(0) * 1000 # ms per sample start = time.time() _ = quantized_model(input_sample, attention_mask=attention_sample) end = time.time() inference_time_after = (end - start) / input_sample.size(0) * 1000 # ms per sample print(f"Validation accuracy after distillation: {val_accuracy:.2f}%") print(f"Model size before quantization: {size_before:.2f} MB") print(f"Model size after quantization: {size_after:.2f} MB") print(f"Inference time before quantization: {inference_time_before:.2f} ms/sample") print(f"Inference time after quantization: {inference_time_after:.2f} ms/sample")
Before Optimization: Accuracy 93%, Size 500MB, Inference 500ms/sample
After Optimization: Accuracy 91.5%, Size 75MB, Inference 120ms/sample
model distillation in NLP?quantization to a model's weights in Python using PyTorch?teacher_outputs = torch.tensor([0.1, 0.9]) student_outputs = torch.tensor([0.1, 0.9]) loss_fn = torch.nn.MSELoss() loss = loss_fn(student_outputs, teacher_outputs) print(loss.item())
AttributeError: 'MyModel' object has no attribute 'quantize'. What is the likely cause?