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等。
单行json19秒读取时间较短,可以直接查看文件需要重新生成新的文件
多行json15秒读取时间较短需要重新生成新的文件,查看不方便,因为第二行全部是全部数据的json串
numpy的loadtxt49秒加载方式较为简单,不用做过多操作需要文件内容的类型一致,否则无法读取,读取时间较慢,性价比不高。
字节文件读取9秒加载速度快需要重新生成文件,而且对于原有字节文件生成的方式要了解,否则无法加载。

相关推荐