bert预训练代码

下面是使用PyTorch实现BERT预训练的代码示例:

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
from torch.utils.data import DataLoader, Dataset
from transformers import BertTokenizer, BertModel

# 载入预训练的BERT模型和词汇表
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertModel.from_pretrained('bert-base-chinese')

# 定义数据集
class MyDataset(Dataset):
    def __init__(self, texts):
        self.texts = texts
    
    def __getitem__(self, index):
        text = self.texts[index]
        tokenized_text = tokenizer(text, padding=True, truncation=True, return_tensors='pt')
        return tokenized_text
    
    def __len__(self):
        return len(self.texts)

# 定义训练参数
batch_size = 32
num_epochs = 10
learning_rate = 1e-4

# 定义模型和优化器
class BertPretrain(nn.Module):
    def __init__(self):
        super(BertPretrain, self).__init__()
        self.bert = model

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids, attention_mask)
        return outputs.last_hidden_state
    
model = BertPretrain()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 载入数据集
dataset = MyDataset(texts)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# 训练模型
for epoch in range(num_epochs):
    for i, batch in enumerate(dataloader):
        input_ids = batch['input_ids'].squeeze(1)
        attention_mask = batch['attention_mask'].squeeze(1)

        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask)
        loss = outputs.mean()
        loss.backward()
        optimizer.step()

        if i % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                  .format(epoch+1, num_epochs, i+1, len(dataloader), loss.item()))

在这个例子中，我们首先使用`BertTokenizer`和`BertModel`从预训练模型中载入BERT模型和词汇表，然后定义了一个数据集`MyDataset`和数据加载器`DataLoader`，并将其用于训练模型。我们还定义了一个`BertPretrain`类，用于实例化BERT模型和定义前向传递。最后，我们使用Adam优化器和均方误差损失训练模型。