python读取大词向量文件
0、前言
我们在工作中经常遇到需要将词向量文件读取到内存,但是正常情况下,我们的单词个数都是数十万个,单词的向量都是几百维,所以导致文件比较大,动辄几个G,在读取文件的时候经常会比较慢,有没有什么办法能够加快读取文件的速度呢,接下来,本人将从如下几种方法,进行速度的对比。
1、文件格式
我们的文件格式是这样,第一行是"单词个数 向量维度",中间用空格分割。以后每行均为"单词\tvalue1 value2 value3 ....."单词和向量之间用"\t"分割,向量之间用空格分割,我们可以取腾讯公开的词向量来进行查看,下面给出示例
100000 768 的 -0.028929112 0.42987955 0.053804845 -0.44394323 0.22613685 -0.23048736 -0.22736746......... 了 -0.19522709 0.5370848 -0.1434914 -0.5097602 0.26118 -0.048514027 -0.30966273 -0.35723355.........
我们这里的实验假定需要将文件读取成data = {‘的‘:[-0.028929112 0.42987955 0.053804845....],‘了‘:[-0.19522709 0.5370848 -0.1434914 -0.5097602...]...}的字典结构。以下给出不同方法的运行时间,由于可能存在代码的问题,所以导致运行时间也会有点出入,发现有问题的小伙伴也可以在评论区评论。
我们这里的测试数据含有10W条的向量数据,所以单词个数为10W,向量维度为768。
2、直接读取
直接读取方式就是从文件中的每一行进行读取,这种方式需要对字符串进行切分,所以总体时间较慢,代码如下
data = {} with open("vocal.vec.100000","r") as f: line = f.readline().strip().split(" ") word_count,dim = int(line[0]),int(line[1]) line = f.readline() while line: line = line.strip().split("\t") if len(line) < 2: line = f.readline() continue word = line[0] vec = [round(float(item), 3) for item in line[1].split(" ")] data[word] = vec line = f.readline()
这种方法最终的运行时间为63秒
3、单行json
单行json是将每一行向量数据存储为一个json串,放置在文件中,首先,我们将原始数据构造成json的数据。
import json # 这一部分和上面的一样 data = {} with open("vocal.vec.100000","r") as f: line = f.readline().strip().split(" ") word_count,dim = int(line[0]),int(line[1]) line = f.readline() while line: line = line.strip().split("\t") if len(line) < 2: line = f.readline() continue word = line[0] vec = [round(float(item), 3) for item in line[1].split(" ")] data[word] = vec line = f.readline() # 构造json print(word_count,dim,sep=" ") for k,v in data.items(): print(json.dumps({k:v})) # 输出到vocal.vec.100000.json文件中
接下来,我们读取json数据
import json data = {} with open("vocal.vec.100000.json","r") as f: line = f.readline().strip().split(" ") word_count,dim = int(line[0]),int(line[1]) line = f.readline() while line: line = line.strip() word_vec = json.loads(line) data.update(word_vec) line = f.readline()
这种方式运行时间是19秒,明显快了很多
4、多行json
多行json是将整个data字典写入到文件,首先我们先生成文件
import json data = {} with open("vocal.vec.100000","r") as f: line = f.readline().strip().split(" ") word_count,dim = int(line[0]),int(line[1]) line = f.readline() while line: line = line.strip().split("\t") if len(line) < 2: line = f.readline() continue word = line[0] vec = [round(float(item), 3) for item in line[1].split(" ")] data[word] = vec line = f.readline() # 生成多行json print(word_count,dim,sep=" ") print(json.dumps(data)) # 输出的文件名字是vocal.vec.100000.json2
我们加载文件
import json data = {} with open("vocal.vec.100000.json2","r") as f: line = f.readline().strip().split("\t") word_count,dim = int(line[0]),int(line[1]) line = f.readline().strip() data = json.loads(line)
最终的时间是15秒,又快了点
5、numpy的loadtxt方法
这种方法利用的numpy的loadtxt方法,由于其有一定的局限性,我们直接给出相应的代码和结果。loadtxt的局限性是文件中所有的数据需要是同一种类型,由于我们的文件数据有int,float和中文文字,所以我们这里只抽取向量的值,即float类型组成文件,加载代码的方式如下
import numpy as np with open("vocal.vec.100000.onlyvec","r") as f: line = f.readline().strip().split(" ") word_count,dim = int(line[0]),int(line[1]) data = np.loadtxt("vocal.vec.100000.onlyvec",dtype=float,skiprows=1)
最终的加载时间是49秒
6、字节文件读取方法
最后,是将数据转变成字节进行读取,首先我们将数据转成字节文件,如下
import struct data = {} with open("vocal.vec.100000.json2","r") as f: line = f.readline().strip().split("\t") word_count,dim = int(line[0]),int(line[1]) line = f.readline().strip() data = json.loads(line) with open("vocal.vec.100000.bin2","wb") as wf: wf.write(struct.pack(‘ii‘,word_count,dim)) for k,v in data.items(): word = k.encode("utf-8") word_len = len(word) wf.write(struct.pack(‘i‘,word_len)) wf.write(word) for vv in v: wf.write(struct.pack("f",vv))
这里我们使用struct方式进行构建,接下来,进行读取
import struct data = {} with open("vocal.vec.100000.bin2","rb") as f: record_struct = struct.Struct("ii") word_count,dim = struct.unpack("ii",f.read(record_struct.size)) for i in range(word_count): record_struct = struct.Struct("i") word_len = struct.unpack("i",f.read(record_struct.size))[0] word = f.read(word_len).decode("utf-8") record_struct = struct.Struct("f"*dim) vec = struct.unpack("f"*dim,f.read(record_struct.size)) data[word] = vec
这种方式最终显示的结果是9秒。
7、总结
我们以一张表格来对这几种方式进行总结
方式 | 时间 | 优点 | 缺点 |
---|---|---|---|
直接读取 | 63秒 | 不用重新修改文件格式,可以直接查看文件 | 读取时间较慢,需要进行一些处理,例如分割字符串,修改float等。 |
单行json | 19秒 | 读取时间较短,可以直接查看文件 | 需要重新生成新的文件 |
多行json | 15秒 | 读取时间较短 | 需要重新生成新的文件,查看不方便,因为第二行全部是全部数据的json串 |
numpy的loadtxt | 49秒 | 加载方式较为简单,不用做过多操作 | 需要文件内容的类型一致,否则无法读取,读取时间较慢,性价比不高。 |
字节文件读取 | 9秒 | 加载速度快 | 需要重新生成文件,而且对于原有字节文件生成的方式要了解,否则无法加载。 |