预训练模型怎么套用?

加载字典和分词工具 实例化分词工具

my_tokenizer = BertTokenizer.from_pretrained(‘…/model/bert-base-chinese’)

加载预训练模型 实例化预训练模型

my_model_pretrained = BertModel.from_pretrained(‘…/model/bert-base-chinese’).to(device)

class MyModel(torch.nn.Module):
def init(self):
super().init()

    # 定义全连接层，将输入特征维度768转换为输出维度2
    self.fc = torch.nn.Linear(768, 2)

def forward(self, input_ids, attention_mask, token_type_ids):
    # 使用预训练模型进行特征提取，设置不训练参数，只提取特征
    # 输入形状: [batch_size, sequence_length]，输出形状: [batch_size, hidden_size]
    with torch.no_grad():
        out = my_model_pretrained(input_ids=input_ids,
                                  attention_mask=attention_mask,
                                  token_type_ids=token_type_ids)

    # 使用全连接层进行下游任务的训练，将特征映射到所需的输出维度
    # 这里取出last_hidden_state的第一个token（通常是[CLS] token）的向量
    # 输入形状: [batch_size, hidden_size]，输出形状: [batch_size, 2]
    out = self.fc(out.last_hidden_state[:, 0])

    # 对输出进行softmax归一化，得到分类概率值
    out = out.softmax(dim=1)

    return out