文本分类实战:fine-tune huggingface Dis

记录使用huggingface transformers 包提供的预训练模型，完成文本分类任务，如何使用 fine-tune Huggingface 的预训练模型 的一个示例记录。

数据集介绍

使用的是IMDB电影评论数据集，这是一个用于二元情感分类的数据集，包含训练集train 和测试集test 各25000条数据，分别包含 正向 positive 和 负面 negative 的评论各占50%，官方的数据是一条评论占一个文本文件记录，我将训练集合测试集中的所有正面和负面的数据分别合并到一个文件中，一行代表一条评论，如下：

可以直接使用命令行下载解压所需要的的数据：

wget http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz
tar -xf aclImdb_v1.tar.gz

Data preprocessing

将所有的数据读入，texts 是一个list ,存放每一条数据，labels 也是一个list，存放的全是 0,1 数字，0 代表negative类型的数据，1代表positive 的数据

def read_imdb_split(split_dir):
    split_dir = Path(split_dir)
    texts = []
    labels = []
    for label_dir in ["pos", "neg"]:
        for text_file in (split_dir/label_dir).iterdir():
            print(text_file)
            with open(text_file,'r') as inp:
              for line in inp:
                texts.append(line.strip())
                labels.append(0 if label_dir == "neg" else 1)
    return texts, labels

train_texts, train_labels = read_imdb_split('train')#25000t条数据切分成neg&pos各一半
test_texts, test_labels = read_imdb_split('test')#25000条数据

DistilBertTokenizerFast

将每句话中的词转换成索引，同时返回模型需要的 atten_mask 等tensor类型的数据

tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')

def data2tensor(sentences,labels):
  input_ids,attention_mask=[],[]
  #input_ids是每个词对应的索引idx ;token_type_ids是对应的0和1，标识是第几个句子；attention_mask是对句子长度做pad
  #input_ids=[22,21,...499] token_type_ids=[0,0,0,0,1,1,1,1] ;attention_mask=[1,1,1,1,1,0,0,0]补零
  for i in range(len(sentences)):
      encoded_dict = tokenizer.encode_plus(
      sentences[i],
      add_special_tokens = False,      # 添加 '[CLS]' 和 '[SEP]'
      max_length = 100,           # 填充 & 截断长度
      pad_to_max_length = True,
      return_tensors = 'pt',         # 返回 pytorch tensors 格式的数据
      )
      input_ids.append(encoded_dict['input_ids'])
      attention_mask.append(encoded_dict['attention_mask'])
  
  input_ids = torch.cat(input_ids, dim=0)#把多个tensor合并到一起
  attention_mask = torch.cat(attention_mask, dim=0)

  input_ids = torch.LongTensor(input_ids)#每个词对应的索引
  attention_mask = torch.LongTensor(attention_mask)#[11100]padding之后的句子
  labels = torch.LongTensor(labels)#所有实例的label对应的索引idx

  return input_ids, attention_mask, labels

数据加载成DataLoader

关于 dataloader 的使用可以看我的这篇博客：
pytorch dataset 类

train_texts, val_texts, train_labels, val_labels = train_test_split(train_texts, train_labels, test_size=.2)

tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')
train_inputs,train_mask,train_labels = data2tensor(train_texts, train_labels)
val_inputs,val_mask,val_labels = data2tensor(val_texts, val_labels)
test_inputs,test_mask,test_labels = data2tensor(test_texts, test_labels)

train_data = Data.TensorDataset(train_inputs, train_mask, train_labels)
train_loader = Data.DataLoader(train_data, batch_size=64, shuffle=True)#遍历train_dataloader 每次返回batch_size条数据
validation_data = Data.TensorDataset(val_inputs, val_mask, val_labels)
val_loader = Data.DataLoader(validation_data, batch_size=64, shuffle=True)

DistilBertForSequenceClassification

使用huggingface API 进行简单的句子分类，设置模型model 参数等：

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

model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased',num_labels=2).to(device)

optim = AdamW(model.parameters(), lr=5e-5)

train and eval

将数据灌入进行训练和验证，model 的返回值查看 DistilBertForSequenceClassification 官方 文档

for epoch in range(2):
  for i,batch in enumerate(train_loader):
    batch = tuple(t.to(device) for t in batch)
    loss=model(input_ids=batch[0], attention_mask=batch[1], labels=batch[2])[0]
    print(loss.item())
    optim.zero_grad()
    loss.backward()
    optim.step()

    if i % 10 == 0:
        eval(model, val_loader)

可以看到官方文档表示如果需要 loss 进行反向传播，取出第一个值即可；如果需要预测值，则取出第二个值；如果需要每层、每个词的隐藏状态，则取出第四个值

eval

训练时候，设置每经过10个batch，在验证集进行验证，在 validation 和 test 时不会传递真是 label 给模型，所以第一个返回值不再是 loss，而是 logits（预测 label 值）
在验证集上计算准确率，保存结果最高的模型参数，用于在新的无标数据上做预测

def save(model):
    # save
    torch.save(model.state_dict(), output_model_file)
    model.config.to_json_file(output_config_file)

def eval(model, val_loader):
  model.eval()
  eval_loss, eval_accuracy, nb_eval_steps = 0, 0, 0
  for batch in val_loader:
    batch = tuple(t.to(device) for t in batch)
    with torch.no_grad():
      logits = model(batch[0], attention_mask=batch[1])[0]#注意跟下面的参数区别，这个地方model.eval()指定是测试，所以没有添加label 
      logits = logits.detach().cpu().numpy()
      label_ids = batch[2].cpu().numpy()
      tmp_eval_accuracy = flat_accuracy(logits, label_ids)
      eval_accuracy += tmp_eval_accuracy
      nb_eval_steps += 1
  print("Validation Accuracy: {}".format(eval_accuracy / nb_eval_steps))
  global best_score
  if best_score < eval_accuracy / nb_eval_steps:
      best_score = eval_accuracy / nb_eval_steps
      save(model)

参考：

英文文本关系抽取
Sequence Classification with IMDb Reviews
BERT Fine-Tuning Tutorial with PyTorch · Chris McCormick (mccormickml.com)