PyTorch预训练Bert模型

本文介绍以下内容：

1. 使用transformers框架做预训练的bert-base模型；
2. 开发平台使用Google的Colab平台，白嫖GPU加速；
3. 使用datasets模块下载IMDB影评数据作为训练数据。

transformers模块简介

transformers框架为Huggingface开源的深度学习框架，支持几乎所有的Transformer架构的预训练模型。使用非常的方便，本文基于此框架，尝试一下预训练模型的使用，简单易用。

本来打算预训练bert-large模型，发现colab上GPU显存不够用，只能使用base版本了。打开colab，并且设置好GPU加速，接下来开始介绍代码。

代码实现

首先安装数据下载模块和transformers包。

!pip install datasets
!pip install transformers

使用datasets下载IMDB数据，返回DatasetDict类型的数据.返回的数据是文本类型，需要进行编码。下面会使用tokenizer进行编码。

from datasets import load_dataset

imdb = load_dataset('imdb')
print(imdb['train'][:3