devbox
从前慢现在也慢
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

transformers训练自己的数据集实战_distilbert-base-uncased

作者：从前慢现在也慢 | 2024-04-01 05:04:28

赞

踩

distilbert-base-uncased
  1. """
  2. 使用IMDb评论进行序列分类
  3. """
  4. #先下载数据
  5. # wget http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz
  6. # tar -xf aclImdb_v1.tar.gz
  7.  
  8. #整理文件
  9. from pathlib import Path
  10.  
  11. def read_imdb_split(split_dir):
  12.     split_dir = Path(split_dir)
  13.     texts = []
  14.     labels = []
  15.     for label_dir in ["pos", "neg"]:
  16.         for text_file in (split_dir/label_dir).iterdir():
  17.             texts.append(text_file.read_text())
  18.             labels.append(0 if label_dir is "neg" else 1)
  19.  
  20.     return texts, labels
  21.  
  22. train_texts, train_labels = read_imdb_split('aclImdb/train')
  23. test_texts, test_labels = read_imdb_split('aclImdb/test')
  24.  
  25. #划分训练集和测试集
  26. from sklearn.model_selection import train_test_split
  27. train_texts, val_texts, train_labels, val_labels = train_test_split(train_texts, train_labels, test_size=.2)
  28.  
  29. #好了，我们已经阅读了数据集。现在让我们讨论令牌化。我们最终将使用预训练的DistilBert训练分类器，因此让我们使用DistilBert标记器
  30. from transformers import DistilBertTokenizerFast
  31. tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')
  32. train_encodings = tokenizer(train_texts, truncation=True, padding=True)
  33. val_encodings = tokenizer(val_texts, truncation=True, padding=True)
  34. test_encodings = tokenizer(test_texts, truncation=True, padding=True)
  35.  
  36. import torch
  37.  
  38. class IMDbDataset(torch.utils.data.Dataset):
  39.     def __init__(self, encodings, labels):
  40.         self.encodings = encodings
  41.         self.labels = labels
  42.  
  43.     def __getitem__(self, idx):
  44.         item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
  45.         item['labels'] = torch.tensor(self.labels[idx])
  46.         return item
  47.  
  48.     def __len__(self):
  49.         return len(self.labels)
  50.  
  51. train_dataset = IMDbDataset(train_encodings, train_labels)
  52. val_dataset = IMDbDataset(val_encodings, val_labels)
  53. test_dataset = IMDbDataset(test_encodings, test_labels)
  54.  
  55. #与培训师进行微调
  56. from transformers import DistilBertForSequenceClassification, Trainer, TrainingArguments
  57.  
  58. training_args = TrainingArguments(
  59.     output_dir='./results',          # output directory
  60.     num_train_epochs=3,              # total number of training epochs
  61.     per_device_train_batch_size=16,  # batch size per device during training
  62.     per_device_eval_batch_size=64,   # batch size for evaluation
  63.     warmup_steps=500,                # number of warmup steps for learning rate scheduler
  64.     weight_decay=0.01,               # strength of weight decay
  65.     logging_dir='./logs',            # directory for storing logs
  66.     logging_steps=10,
  67. )
  68.  
  69. model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")
  70.  
  71. trainer = Trainer(
  72.     model=model,                         # the instantiated 
    声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/从前慢现在也慢/article/detail/347526
    推荐阅读
    相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。