更多基础解释见BERT分类构造
代码及注释:
class DistilBertForSequenceClassification(DistilBertPreTrainedModel):
def __init__(self, config: PretrainedConfig):
super().__init__(config)
self.num_labels = config.num_labels
self.config = config
self.distilbert = DistilBertModel(config)
self.pre_classifier = nn.Linear(config.dim, config.dim)
self.classifier = nn.Linear(config.dim, config.num_labels)
self.dropout = nn.Dropout(config.seq_classif_dropout)
self.post_init()
def get_position_embeddings(self) -> nn.Embedding:
return self.distilbert.get_position_embeddings()
def resize_position_embeddings(self, new_num_position_embeddings: int):
self.distilbert.resize_position_embeddings(new_num_position_embeddings)
def forward(
self,
input_ids: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
head_mask: Optional[torch.Tensor] = None,
inputs_embeds: Optional[torch.Tensor] = None,
labels: Optional[torch.LongTensor] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> Union[SequenceClassifierOutput, Tuple[torch.Tensor, ...]]:
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
distilbert_output = self.distilbert(
input_ids=input_ids,
attention_mask=attention_mask,
head_mask=head_mask,
inputs_embeds=inputs_embeds,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
"""
第一个元素(索引为0)是模型最后一层的隐藏状态。
它具有形状(batch_size, sequence_length, hidden_dimension),
表示对于每个批次中的样本,每个token都有一个dim维的向量表示。
"""
hidden_state = distilbert_output[0] # (bs, seq_len, dim)
"""
从上述隐藏状态中,选取每个样本的第一个token(通常为[CLS]标记)的向量作为池化输出。
pooled_output变成了形状为(batch_size, hidden_dimension)的向量,每个样本一个固定长度的向量表示
"""
pooled_output = hidden_state[:, 0] # (bs, dim)
"""
对特征进行线性变换.
过渡层,后跟激活函数,增强模型的表达能力、帮助模型学习到更高级别的抽象特征?
"""
pooled_output = self.pre_classifier(pooled_output) # (bs, dim)
"""
应用ReLU激活函数到pre_classifier的输出上。
引入非线性,帮助模型学习更复杂的模式,提高其表达能力。
"""
pooled_output = nn.ReLU()(pooled_output) # (bs, dim)
"""应用Dropout层,随机“丢弃”一部分神经元的输出以减少过拟合。"""
pooled_output = self.dropout(pooled_output) # (bs, dim)
"""
该层负责将高维特征映射到类别空间,
输出形状为(batch_size, num_labels)的张量,即每个样本对于每个类别的logits(未经过softmax的预测分数)。
这些logits随后可用于计算分类损失或直接进行预测。
"""
logits = self.classifier(pooled_output) # (bs, num_labels)
loss = None
if labels is not None:
if self.config.problem_type is None:
if self.num_labels == 1:
self.config.problem_type = "regression"
elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
self.config.problem_type = "single_label_classification"
else:
self.config.problem_type = "multi_label_classification"
if self.config.problem_type == "regression":
loss_fct = MSELoss()
if self.num_labels == 1:
loss = loss_fct(logits.squeeze(), labels.squeeze())
else:
loss = loss_fct(logits, labels)
elif self.config.problem_type == "single_label_classification":
loss_fct = CrossEntropyLoss()
loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
elif self.config.problem_type == "multi_label_classification":
loss_fct = BCEWithLogitsLoss()
loss = loss_fct(logits, labels)
if not return_dict:
output = (logits,) + distilbert_output[1:]
return ((loss,) + output) if loss is not None else output
return SequenceClassifierOutput(
loss=loss,
logits=logits,
hidden_states=distilbert_output.hidden_states,
attentions=distilbert_output.attentions,
)