python读取大词向量文件

目录

• 0、前言
• 1、文件格式
• 2、直接读取
• 3、单行json
• 4、多行json
• 5、numpy的loadtxt方法
• 6、字节文件读取方法
• 7、文件加载
• 8、总结

0、前言

我们在工作中经常遇到需要将词向量文件读取到内存，但是正常情况下，我们的单词个数都是数十万个，单词的向量都是几百维，所以导致文件比较大，动辄几个G，在读取文件的时候经常会比较慢，有没有什么办法能够加快读取文件的速度呢，接下来，本人将从如下几种方法，进行速度的对比。

1、文件格式

我们的文件格式是这样，第一行是"单词个数 向量维度"，中间用空格分割。以后每行均为"单词 value1 value2 value3 ....."单词和向量之间用" "分割，向量之间用空格分割，我们可以取腾讯公开的词向量来进行查看，下面给出示例

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("	")
    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("	")
    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("	")
    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("	")
  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("	")
  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秒	加载速度快	需要重新生成文件，而且对于原有字节文件生成的方式要了解，否则无法加载。