基于Huggingface使用BERT进行文本分类的fine-tuning

基于Huggingface使用BERT进行文本分类的fine-tuning

随着BERT大火之后，很多BERT的变种，这里借用Huggingface工具来简单实现一个文本分类，从而进一步通过Huggingface来认识BERT的工程上的实现方法。

1、load data

train_df = pd.read_csv('../data/train.tsv',delimiter='\t',names=['text','label'])

print(train_df.shape)

train_df.head()

sentences = list(train_df['text'])

targets =train_df['label'].values

2、token encodding

#如果token要封装到自定义model类中的话，则需要指定max_len

tokenizer=BertTokenizer.from_pretrained('bert-base-uncased')

max_length=32

sentences_tokened=tokenizer(sentences,padding=True,truncation=True,max_length=max_length,return_tensors='pt')

targets=torch.tensor(targets)

3、encoding data

# from torchvision import transforms,datasets

from torch.utils.data import Dataset,DataLoader,random_split

class DataToDataset(Dataset):

    def __init__(self,encoding,labels):

        self.encoding=encoding

        self.labels=labels

    def __len__(self):

        return len(self.labels)

    def __getitem__(self,index):

        return self.encoding['input_ids'][index],self.encoding['attention_mask'][index],self.labels[index]

#封装数据

datasets=DataToDataset(sentences_tokened,targets)

train_size=int(len(datasets)*0.8)

test_size=len(datasets)-train_size

print([train_size,test_size])

train_dataset,val_dataset=random_split(dataset=datasets,lengths=[train_size,test_size])

BATCH_SIZE=64

#这里的num_workers要大于0

train_loader=DataLoader(dataset=train_dataset,batch_size=BATCH_SIZE,shuffle=True,num_workers=5)

val_loader=DataLoader(dataset=val_dataset,batch_size=BATCH_SIZE,shuffle=True,num_workers=5)#

4、create model

class BertTextClassficationModel(nn.Module):

    def __init__(self):

        super(BertTextClassficationModel,self).__init__()

        self.bert=BertModel.from_pretrained('bert-base-uncased')

        self.dense=nn.Linear(768,2)  #768 input, 2 output

    def forward(self,ids,mask):

        out,_=self.bert(input_ids=ids,attention_mask=mask)

        out=self.dense(out[:,0,:])

        return out

mymodel=BertTextClassficationModel()

#获取gpu和cpu的设备信息

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("device=",device)

if torch.cuda.device_count()>1:

    print("Let's use ",torch.cuda.device_count(),"GPUs!")

    mymodel=nn.DataParallel(mymodel)

mymodel.to(device)

5、train model

loss_func=nn.CrossEntropyLoss()

optimizer=optim.Adam(mymodel.parameters(),lr=0.0001)

from sklearn.metrics import accuracy_score

def flat_accuracy(preds,labels):

    pred_flat=np.argmax(preds,axis=1).flatten()

    labels_flat=labels.flatten()

    return accuracy_score(labels_flat,pred_flat)

epochs=3

for epoch in range(epochs):

    train_loss = 0.0

    train_acc=0.0

    for i,data in enumerate(train_loader):

        input_ids,attention_mask,labels=[elem.to(device) for elem in data]

        #优化器置零

        optimizer.zero_grad()

        #得到模型的结果

        out=mymodel(input_ids,attention_mask)

        #计算误差

        loss=loss_func(out,labels)

        train_loss += loss.item()

        #误差反向传播

        loss.backward()

        #更新模型参数

        optimizer.step()

        #计算acc

        out=out.detach().numpy()

        labels=labels.detach().numpy()

        train_acc+=flat_accuracy(out,labels)

    print("train %d/%d epochs Loss:%f, Acc:%f" %(epoch,epochs,train_loss/(i+1),train_acc/(i+1)))

6、evaluate

print("evaluate...")

val_loss=0

val_acc=0

mymodel.eval()

for j,batch in enumerate(val_loader):

    val_input_ids,val_attention_mask,val_labels=[elem.to(device) for elem in batch]

    with torch.no_grad():

        pred=mymodel(val_input_ids,val_attention_mask)

        val_loss+=loss_func(pred,val_labels)

        pred=pred.detach().cpu().numpy()

        val_labels=val_labels.detach().cpu().numpy()

        val_acc+=flat_accuracy(pred,val_labels)

print("evaluate loss:%d, Acc:%d" %(val_loss/len(val_loader),val_acc/len(val_loader)))