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、文件加载
这里再简单介绍一下本人利用字节进行加载后遇到的问题。由于python的变量存储方式，将一个int整型存储到内存中会用到28个字节，这个和我们的日常认知不同，我们日常认知一个int占用4个字节。可以看下下面的示例：
import sys
a = 0
b = 1
c = 1.0
print(sys.getsizeof(a)) # 24
print(sys.getsizeof(b)) # 28
print(sys.getsizeof(c)) # 24
d = []
e = set()
f = {}
print(sys.getsizeof(d)) # 72
print(sys.getsizeof(e)) # 232
print(sys.getsizeof(f)) # 248
由上面可以看到，我们一个float星占用了24个字节，当我们构建data = {word1:vec1,word2:vec2...}的时候，这个data字典会占用非常大的内存，所以我们需要解决这个问题，我这边解决的办法是用字节作为data的Key和value，在加载数据文件时（这里加载字节文件）：
import struct
self.word2vec = {}
self.word_count = 0
self.dim = 0
with open(model_path,"rb") as f:
record_struct = struct.Struct("ii")
self.word_count,self.dim = struct.unpack("ii",f.read(record_struct.size))
for i in range(self.word_count):
record_struct = struct.Struct("i")
word_len = struct.unpack("i",f.read(record_struct.size))[0]
word = f.read(word_len)
record_struct = struct.Struct("f"*self.dim)
vec = f.read(record_struct.size)
self.word2vec[word] = vec
这样的内存占用就少了，我试了下，内存直接从17%降到了1.9%。效果还是很明显的。在使用的时候，我们将字节在转成相应的类型就可以了。
8、总结
我们以一张表格来对这几种方式进行总结
方式 | 时间 | 优点 | 缺点
- | - | - | -
直接读取 | 63秒 | 不用重新修改文件格式，可以直接查看文件 | 读取时间较慢，需要进行一些处理，例如分割字符串，修改float等。
单行json | 19秒 | 读取时间较短，可以直接查看文件 | 需要重新生成新的文件
多行json | 15秒 | 读取时间较短 | 需要重新生成新的文件，查看不方便，因为第二行全部是全部数据的json串
numpy的loadtxt | 49秒 | 加载方式较为简单，不用做过多操作 | 需要文件内容的类型一致，否则无法读取，读取时间较慢，性价比不高。
字节文件读取 | 9秒 | 加载速度快 | 需要重新生成文件，而且对于原有字节文件生成的方式要了解，否则无法加载。
