HuggingFace | 如何下载预训练模型

发布时间 2023-07-05 12:39:07作者: 张Zong在修行

本例我们在Linux上进行下载,下载的模型是bert-base-uncased

下载网址为:https://www.huggingface.co/bert-base-uncased

huggingface的transformers框架,囊括了BERT、GPT、GPT2、ToBERTa、T5等众多模型,同时支持pytorch和tensorflow 2,代码非常规范,使用也非常简单,但是模型使用的时候,要从他们的服务器上去下载模型,那么有没有办法,把这些预训练模型下载好,在使用时指定使用这些模型呢?答案是肯定的。

总览

总体是,将所需要的预训练模型、词典等文件下载至本地文件夹中 ,然后加载的时候model_name_or_path参数指向文件的路径即可。

下载方法

说明:我这三种下载方法都亲自尝试过,若你不成功请重复尝试,如果卡住,把下载一半的文件删掉,在重启一下重新尝试。

方法一:代码直接下载

我们在第一次执行BertTokenizer BertModel中的某个模型的from_pretrained函数的时候,将会自动下载预训练模型的相关文件。Linux中默认下载到~/.cache/huggingface/transformers中。

代码如下:

from transformers import BertTokenizer, BertModel
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertModel.from_pretrained("bert-base-uncased")

# ...一系列操作
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt')
output = model(**encoded_input)
print(output)

方法二:网址下载

首先打开网址:https://huggingface.co/models 这个网址是huggingface/transformers支持的所有模型,目前一直在增长。搜索bert-base-uncased(其他的模型类似,比如gpt2等),并点击进去。

点击页面中的Files and versions,可以看到模型的所有文件。

通常我们需要保存的是三个文件及一些额外的文件,第一个是配置文件;config.json。第二个是词典文件,vocab.json。第三个是预训练模型文件,如果你使用pytorch则保存pytorch_model.bin文件,如果你使用tensorflow 2,则保存tf_model.h5

额外的文件,指的是tokenizer.jsontokenizer_config.json等,这几类是tokenizer需要使用的文件,如果出现的话,也需要保存下来。没有的话,就不必在意。如果不确定哪些需要下,哪些不需要的话,可以把上图中类似的文件或全部文件全都下载下来。

下载到本地文件夹bert-base-uncased,同时这些名称务必注意保证与图中的名称完全一致。然后从本地上传至自己的服务器想要存储的位置。

文件结构:

|- bertLearning 项目目录
|---->bert-base-uncased  下载好的bert-base-uncased预训练模型的文件,放在名为bert-base-uncased文件夹下
|---->run.py  加载执行预训练模型的文件

run.py代码如下:

from transformers import BertTokenizer, BertModel
tokenizer = BertTokenizer.from_pretrained("./bert-base-uncased")
model = BertModel.from_pretrained("./bert-base-uncased")

# ...一系列操作
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt')
output = model(**encoded_input)
print(output)

我们可以看出from_pretrained函数的参数是一个字符串类型的路径,例子中使用的路径是相对路径。意思是运行的这个py文件的文件夹中的bert-base-uncased文件。然后,就会自动在这个文件夹下去寻找相应的文件去加载tokenizer和model。

方法三:添加cache_dir参数

我们在第一次执行BertTokenizer BertModel中的某个模型的from_pretrained函数的时候,将会自动下载预训练模型的相关文件。但是我们还在函数中增加了cache_dir参数,这个参数的作用是你要下载到自己想要指定的一个文件夹下。

代码如下:

from transformers import BertTokenizer, BertModel
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased",cache_dir="/home/user/path/to/bertLearning/cache/")
model = BertModel.from_pretrained("bert-base-uncased",cache_dir="/home/user/path/to/bertLearning/cache/")

# ...一系列操作
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt')
output = model(**encoded_input)
print(output)

下载完,我们创建的cache文件夹中的预训练模型的文件如下:

总结:

  • 千万不要把三个方法混着用,可能你最后都不知道到底哪个文件夹加载起了作用。(在全都下载的情况下,可能存在优先级查找)
  • 上面的三种方法亲身实验,如果不成功,删除和重启从新下载文件。