使用DistilBERT 蒸餾類(lèi) BERT 模型的代碼實(shí)現(xiàn)

from transformers import AutoModelForMaskedLM
roberta = AutoModelForMaskedLM.from_pretrained("roberta-large")
print(roberta)

from typing import Anyfrom transformers import AutoModelForMaskedLM
roberta = AutoModelForMaskedLM.from_pretrained("roberta-large")
def visualize_children(  object : Any,  level : int = 0,) -> None:  """  Prints the children of (object) and their children too, if there are any.  Uses the current depth (level) to print things in a ordonnate manner.  """  print(f"{'   ' * level}{level}- {type(object).__name__}")  try:      for child in object.children():          visualize_children(child, level + 1)  except:      pass
visualize_children(roberta)

from transformers.models.roberta.modeling_roberta import RobertaPreTrainedModel, RobertaConfig
def distill_roberta(  teacher_model : RobertaPreTrainedModel,) -> RobertaPreTrainedModel:  """  Distilates a RoBERTa (teacher_model) like would DistilBERT for a BERT model.  The student model has the same configuration, except for the number of hidden layers, which is // by 2.  The student layers are initilized by copying one out of two layers of the teacher, starting with layer 0.  The head of the teacher is also copied.  """  # Get teacher configuration as a dictionnary  configuration = teacher_model.config.to_dict()  # Half the number of hidden layer  configuration['num_hidden_layers'] //= 2  # Convert the dictionnary to the student configuration  configuration = RobertaConfig.from_dict(configuration)  # Create uninitialized student model  student_model = type(teacher_model)(configuration)  # Initialize the student's weights  distill_roberta_weights(teacher=teacher_model, student=student_model)  # Return the student model  return student_model

from transformers.models.roberta.modeling_roberta import RobertaEncoder, RobertaModelfrom torch.nn import Module
def distill_roberta_weights(  teacher : Module,  student : Module,) -> None:  """  Recursively copies the weights of the (teacher) to the (student).  This function is meant to be first called on a RobertaFor... model, but is then called on every children of that model recursively.  The only part that's not fully copied is the encoder, of which only half is copied.  """  # If the part is an entire RoBERTa model or a RobertaFor..., unpack and iterate  if isinstance(teacher, RobertaModel) or type(teacher).__name__.startswith('RobertaFor'):      for teacher_part, student_part in zip(teacher.children(), student.children()):          distill_roberta_weights(teacher_part, student_part)  # Else if the part is an encoder, copy one out of every layer  elif isinstance(teacher, RobertaEncoder):          teacher_encoding_layers = [layer for layer in next(teacher.children())]          student_encoding_layers = [layer for layer in next(student.children())]          for i in range(len(student_encoding_layers)):              student_encoding_layers[i].load_state_dict(teacher_encoding_layers[2*i].state_dict())  # Else the part is a head or something else, copy the state_dict  else:      student.load_state_dict(teacher.state_dict())

from torch import Tensor
def get_logits(  model : RobertaPreTrainedModel,  input_ids : Tensor,  attention_mask : Tensor,) -> Tensor:  """  Given a RoBERTa (model) for classification and the couple of (input_ids) and (attention_mask),  returns the logits corresponding to the prediction.  """  return model.classifier(      model.roberta(input_ids, attention_mask)[0]  )

import torchfrom torch.nn import CrossEntropyLoss, CosineEmbeddingLoss
def distillation_loss(  teacher_logits : Tensor,  student_logits : Tensor,  labels : Tensor,  temperature : float = 1.0,) -> Tensor:  """  The distillation loss for distilating a BERT-like model.  The loss takes the (teacher_logits), (student_logits) and (labels) for various losses.  The (temperature) can be given, otherwise it's set to 1 by default.  """  # Temperature and sotfmax  student_logits, teacher_logits = (student_logits / temperature).softmax(1), (teacher_logits / temperature).softmax(1)  # Classification loss (problem-specific loss)  loss = CrossEntropyLoss()(student_logits, labels)  # CrossEntropy teacher-student loss  loss = loss + CrossEntropyLoss()(student_logits, teacher_logits)  # Cosine loss  loss = loss + CosineEmbeddingLoss()(teacher_logits, student_logits, torch.ones(teacher_logits.size()[0]))  # Average the loss and return it  loss = loss / 3  return loss