python讀取大詞向量文件

目錄

• 0、前言
• 1、文件格式
• 2、直接讀取
• 3、單行json
• 4、多行json
• 5、numpy的loadtxt方法
• 6、字節文件讀取方法
• 7、文件加載
• 8、總結

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秒	加載速度快	需要重新生成文件，而且對於原有字節文件生成的方式要了解，否則無法加載。