用tensflow实现word2vec
Tue, Jun 13, 2017 in Interesting Python
word2vec是一个google提出的，用于学习词嵌入的模型，模型可以学习每个短语的语义，并对其进行向量化
实现参考了一部分github tensoflow demo 的代码，最后的结果如下，可以看到算法还是找出一些语义相近的词的:
这个版本的word2vec已经开源在我的github
word2vec
word2vec实现了从单词到词向量的映射，语义相近的词的向量会很接近，比如上图中的beautiful和gorgeous，morning和night，guys和people，之前也在很多场景下用过word2vec，但是一直以来word2vec对于我来说是一个黑盒子，只要给它一些文章，最后它就能学习单词对应的空间向量，然后这些向量可以用来做一些更加深入的应用，比如作为深度学习的输入，进行文本类别或者情感分类，或者是寻找近同义词之类的，之前实习过的某个公司用word2vec判断公司名字的类别，效果不错。
网上描述word2vec的训练方法的文章有很多，但是word2vec的训练方法大值只有两种，CBOW和skip-gram，google的tensorflow demo给出了skip-gram的方法，这里我实现的是CBOW的方法：
the code
from __future__ import print_function
import nltk
from collections import Counter
import tensorflow
import keras
import numpy as np
from keras import backend as K
from keras.engine.topology import Layer
import random
import os
import sys
import time
NUM_WORDS = 40000
C_WINDOW = 1
EMBEDDING_DIM = 128
batch_size = 512
neg_size = 64
text = open('chat.txt',encoding='utf8').read()
words = nltk.tokenize.word_tokenize(text)
counter_words = Counter(words)
english_punctuations = [',', '.', ':', ';', '?', '(', ')', '[', ']', '!', '@', '#', '%', '$', '*']
english_punctuations = dict(zip(english_punctuations,range(len(english_punctuations))))
words = [i for i in words if i not in english_punctuations]
wordseq = sorted(counter_words.items(),key=lambda x:x[1])[-NUM_WORDS + 1:]
word2ind = {}
ind2word = {}
for ind,val in enumerate(wordseq):
ind += 1
word,freq = val
word2ind[word] = ind
ind2word[ind] = word
word2ind['unk'] = 0
ind2word[0] = 'unk'
indseq = [word2ind[i] if i in word2ind else 0 for i in words]\
indseq = np.asarray(indseq,dtype=np.int)
xs = []
ys = []
for i,val in enumerate(indseq):
if i < C_WINDOW or i >= len(indseq) - C_WINDOW:
continue
xs.append(np.concatenate((indseq[i - C_WINDOW:i],indseq[i + 1:i + C_WINDOW + 1] ),axis=0))
ys.append(val)
xs = np.asarray(xs)
ys = np.asarray(ys)
print(xs.shape,ys.shape)
class Dataset():
def __init__(self,data,label):
self._index_in_epoch = 0
self._epochs_completed = 0
self._data = data
self._label = label
assert(data.shape[0] == label.shape[0])
self._num_examples = data.shape[0]
pass
@property
def data(self):
return self._data
@property
def label(self):
return self._label
def next_batch(self,batch_size,shuffle = True):
start = self._index_in_epoch
if start == 0 and self._epochs_completed == 0:
idx = np.arange(0, self._num_examples) # get all possible indexes
np.random.shuffle(idx) # shuffle indexe
self._data = self.data[idx] # get list of `num` random samples
self._label = self.label[idx]
# go to the next batch
if start + batch_size > self._num_examples:
self._epochs_completed += 1
rest_num_examples = self._num_examples - start
data_rest_part = self.data[start:self._num_examples]
label_rest_part = self.label[start:self._num_examples]
idx0 = np.arange(0, self._num_examples) # get all possible indexes
np.random.shuffle(idx0) # shuffle indexes
self._data = self.data[idx0] # get list of `num` random samples
self._label = self.label[idx0]
start = 0
self._index_in_epoch = batch_size - rest_num_examples #avoid the case where the #sample != integar times of batch_size
end = self._index_in_epoch
data_new_part = self._data[start:end]
label_new_part = self._label[start:end]
return np.concatenate((data_rest_part, data_new_part), axis=0),np.concatenate((label_rest_part, label_new_part), axis=0)
else:
self._index_in_epoch += batch_size
end = self._index_in_epoch
return self._data[start:end],self._label[start:end]
dset = Dataset(xs,ys)
class Word2vecLayer(Layer):
def __init__(self, word_num,embedding_size, **kwargs):
self.word_num = word_num
self.embedding_size = embedding_size
super(Word2vecLayer, self).__init__(**kwargs)
def build(self, input_shape):
# Create a trainable weight variable for this layer.
self.weight_matrix = self.add_weight(shape=(self.word_num, self.embedding_size),
initializer='uniform',name='weight_matrix',
trainable=True)
super(Word2vecLayer, self).build(input_shape) # Be sure to call this somewhere!
def call(self, inputs):
if K.dtype(inputs) != 'int32':
inputs = K.cast(inputs, 'int32')
out = K.gather(self.weight_matrix, inputs)
out = K.sum(out,axis=1)
return out
def compute_output_shape(self, input_shape):
return (input_shape[0], self.embedding_size)
x_input = tf.placeholder(tf.int32,shape=(None,C_WINDOW * 2))
y_input = tf.placeholder(tf.int32,shape=(None,1))
word2veclayer = Word2vecLayer(NUM_WORDS,EMBEDDING_DIM)
model = word2veclayer(x_input)
show_word_indexs = random.sample([word2ind[i[0]] for i in wordseq[-400:-200] if len(i[0]) > 1],20)
nce_weight = tf.Variable(tf.truncated_normal([NUM_WORDS,EMBEDDING_DIM],-1,1))
nce_bias = tf.Variable(tf.zeros([NUM_WORDS,]))
loss = tf.reduce_mean(tf.nn.nce_loss(nce_weight, nce_bias,
labels=y_input,
inputs=model,
num_sampled=neg_size,
num_classes=NUM_WORDS))
norm = tf.sqrt(tf.reduce_sum(tf.square(word2veclayer.weight_matrix), 1, keep_dims=True))
normalized_embeddings = word2veclayer.weight_matrix / norm
valid_words = tf.constant(show_word_indexs)
valid_embeddings = tf.nn.embedding_lookup(normalized_embeddings, valid_words)
similarity = tf.matmul(valid_embeddings, tf.transpose(normalized_embeddings))
optimizer = tf.train.AdagradOptimizer(1.0).minimize(loss)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
class ProgressBar():
def __init__(self,worksum,info="",auto_display=True):
self.worksum = worksum
self.info = info
self.finishsum = 0
self.auto_display = auto_display
def startjob(self):
self.begin_time = time.time()
def complete(self,num):
self.gaptime = time.time() - self.begin_time
self.finishsum += num
if self.auto_display == True:
self.display_progress_bar()
def display_progress_bar(self):
percent = self.finishsum * 100 / self.worksum
eta_time = self.gaptime * 100 / (percent + 0.001) - self.gaptime
strprogress = "[" + "=" * int(percent // 2) + ">" + "-" * int(50 - percent // 2) + "]"
str_log = ("%s %.2f %% %s %s/%s \t used:%ds eta:%d s" % (self.info,percent,strprogress,self.finishsum,self.worksum,self.gaptime,eta_time))
sys.stdout.write('\r' + str_log)
for i in range(5):
pb = ProgressBar(worksum=len(indseq),info='itertion %s' % (i + 1))
pb.startjob()
for j in range(len(indseq) // batch_size):
batch_x,batch_y = dset.next_batch(batch_size)
batch_y = batch_y.reshape(-1,1)
sess.run(optimizer,feed_dict={x_input:batch_x,y_input:batch_y})
pb.complete(batch_size)
if i % 1 == 0:
sim = similarity.eval(session=sess)
for i in range(len(show_word_indexs)):
valid_word = ind2word[show_word_indexs[i]]
top_k = 8 # number of nearest neighbors
nearest = (-sim[i, :]).argsort()[1:top_k+1]
log = 'Nearest to %s:' % valid_word
for k in range(top_k):
close_word = ind2word[nearest[k]]
log = '%s %s,' % (log, close_word)
print(log)
训练所用到的推特语料在这里 可以找到，大概20Mb左右，tensorflow训练一遍语料需要大概150秒，这是在内存32Gb，装有1080ti显卡的机器上的速度，从代码中也可以看出一共训练了5轮。
CBOW
word2vec的论文中在输出层用的树形结构减少内存的trick我们是没有用到的，这里实现的仅仅是一个最简单的word2vec的模型
我们的CBOW模型也由一个输入层，一个projection层和一个output层组成模型的核心有两处：
projection layer
class Word2vecLayer(Layer):
def __init__(self, word_num,embedding_size, **kwargs):
self.word_num = word_num
self.embedding_size = embedding_size
super(Word2vecLayer, self).__init__(**kwargs)
def build(self, input_shape):
# Create a trainable weight variable for this layer.
self.weight_matrix = self.add_weight(shape=(self.word_num, self.embedding_size),
initializer='uniform',name='weight_matrix',
trainable=True)
super(Word2vecLayer, self).build(input_shape) # Be sure to call this somewhere!
def call(self, inputs):
if K.dtype(inputs) != 'int32':
inputs = K.cast(inputs, 'int32')
out = K.gather(self.weight_matrix, inputs)
out = K.sum(out,axis=1)
return out
def compute_output_shape(self, input_shape):
return (input_shape[0], self.embedding_size)
这个layer就是中间的，projection层，逻辑就是将某个词之前之后WINDOW中的词对应的词向量相加然后输出,本来是写给keras用的，但是后来发现keras中没有negative sample error的实现，就转用tensorflow写了，但是tensorflow中可以直接适用keras的模型，所以模型直接使用了这一层，而模型也只有这一层。
word2veclayer = Word2vecLayer(NUM_WORDS,EMBEDDING_DIM)
model = word2veclayer(x_input)
negative sample error
negative sample error 可以很好的避免softmax在维度较大的情况下计算量的问题，在这里用于加快softmax层的计算速度，NCE的思想是，对于每一个正样本，随机选择N个负样本，一共N + 1个样本计算loss，替代原先的计算所有loss的方法，当样本数（在word2vec中是单词数）增加得比较剧烈时，negative sample error 可以做到对训练速度影响不大。
训练
我首先选择了《三国演义》进行训练，我下载得《三国演义》版本去掉无用的符号以后一共是28万字，这对于word2vec来说是很少的量，我一共训练了5轮，每轮15秒左右，每轮结果是这样的：
itertion 1 99.98 % [=================================================>-] 283904/283961 used:15s eta:0 sNearest to 蜀兵: 触先言, 却引兵, 拥而来, 国宝, 惭恨, 最喜, 利害, 托以,
Nearest to 后人: 再三, 折二将, 一箭, 劫寨, 门外, 屯田, 尘头, 之托,
Nearest to 上马: 出马, 接肃入, 危急, 必死于, 大怒, 擒下, 提刀, 舞罢,
Nearest to 次日: 造饭, 以练, 延, 死活, 孔明自, 孟达, 驱得动, 长沙,
Nearest to 众将: 未能, 辞职, 李严, 尽迁, 多问, 旧交, 恩, 无限,
Nearest to 一面: 从事, 飞报, 安能, 关上, 动, 率兵, 劳军, 废,
Nearest to 太守: 司马昭, 刀斧手, 所, 风雪, 方到, 南安, 此事, 相待,
Nearest to 不能: 无不, 三人, 虽然, 必来, 往求, 于桓, 重围, 炭,
Nearest to 引兵: 之兵, 拜服, 下山, 三人, 地脉, 飞马, 领, 领兵,
Nearest to 人马: 泠, 人到, 兴兵, 我军, 聚集, 聚, 信, 专等,
Nearest to 关公: 其计, 计, 街亭, 关兴, 人到, 举刀, 料, 安有,
Nearest to 一人: 之间, 魏能, 谏, 诈修, 曹寨, 召公瑾, 射中, 手,
Nearest to 玄德曰: 出榜, 轻敌, 孔明惊, 尽打, 刘景升, 离, 银冶洞, 痛恨,
Nearest to 袁绍: 山上, 西, 朝见, 尚, 从事, 审配, 当阳, 可以,
Nearest to 都督: 刺, 哥哥, 守寨, 放箭, 张达, 禁从, 之学, 回魏,
Nearest to 却说: 惊, 参军, 公既为, 做作, 司马昭之, 命玄, 屯田, 赍,
Nearest to 东吴: 大兵, 成功, 尽力, 回见, 臣安得, 傅士仁, 乘, 可用,
Nearest to 刘备: 司马隽, 若见, 出征, 义结, 元绍汉, 玄德背, 入府, 亡年,
Nearest to 赵云: 赶至, 掘, 引见, 褚上, 床头, 自引兵, 可下, 由,
Nearest to 二人: 殿, 称王, 本寨, 张鲁之, 答礼, 伏后, 满面, 不得已,
itertion 3 99.98 % [=================================================>-] 283904/283961 used:15s eta:0 sNearest to 蜀兵: 李所杀, 使于, 各怀, 同破, 绍大怒, 承惊, 雕, 幸臣,
Nearest to 后人: 江湖, 御前, 定计, 平怒, 彼之怒, 松, 云弃, 鲜明,
Nearest to 上马: 西凉兵, 石二寨, 飞转, 跳跃, 直打, 军中无戏言, 广陵, 谗,
Nearest to 次日: 一切, 岂知, 至家, 王气, 云领, 王远置, 来冲, 免忧,
Nearest to 众将: 谭尸, 妄称, 分别, 营中掘, 正慌急, 声势, 沿城, 邓芝,
Nearest to 一面: 入赘, 宫殿, 恐是, 炎见孚, 空阔, 恨不得, 何来, 张鲁见,
Nearest to 太守: 二者, 志, 何故, 那马, 马刺, 恪自, 君辱臣死, 周出,
Nearest to 不能: 附贼, 文安, 一武, 今公, 谋害, 邀请, 欲哭无泪, 危急,
Nearest to 引兵: 借得, 下山, 飞马, 敢, 冲出重围, 章武, 差官, 乃举,
Nearest to 人马: 仁义, 保守, 此二快, 屯兵, 差, 法正谓, 军杀, 两头,
Nearest to 关公: 容彼自, 相哭, 飞至, 十五日, 玄德再, 杨暨入, 勒住, 或言,
Nearest to 一人: 获一人, 一人出, 权取, 推尊, 英气, 之间, 喜汝有, 上车,
Nearest to 玄德曰: 长戟, 必须, 竭力, 别计, 歇, 国太问, 轻敌, 不留,
Nearest to 袁绍: 云乃拔, 药, 五十万, 过度, 话分, 涿, 天幸, 三通,
Nearest to 都督: 救于, 魏屯, 青旗军, 海若, 看, 五虎, 五分, 如玉,
Nearest to 却说: 水火之中, 如昔, 公既为, 拈阄, 攀住, 卓擎, 恭俭节, 催兵,
Nearest to 东吴: 西接, 乃亦布, 阳曲, 城可得, 终日, 超视, 斗得, 引二子,
Nearest to 刘备: 庶民, 一万五千, 为言, 独行, 入府, 范请, 依得, 鹏,
Nearest to 赵云: 复命, 各藏, 孟优来, 不苦, 武骑, 未迟, 睿知, 而见,
Nearest to 二人: 称王, 伏法, 频频, 明此意, 北山, 答礼, 糜芳, 世所,
itertion 5 99.98 % [=================================================>-] 283904/283961 used:15s eta:0 sNearest to 蜀兵: 司马懿留, 拥上, 投雍寨, 反乱, 完卵, 李所杀, 劫迁, 梁上,
Nearest to 后人: 分后, 刺弘, 偏护, 申奏魏, 当休, 暗, 犹有, 经典,
Nearest to 上马: 册封, 张大, 休慌, 我荐, 谕以, 西凉兵, 迁劫汉帝, 跳跃,
Nearest to 次日: 后果, 岂知, 吕布来, 说否, 名可, 尽带, 二日, 李劫,
Nearest to 众将: 错动, 丁等, 石苞为, 我们, 以纪, 夸美, 力困, 张先,
Nearest to 一面: 感晋臣, 入赘, 拂袖, 盗, 何来, 飞罪, 自纵, 帆,
Nearest to 太守: 看此, 不难, 刘琦共投, 包原, 吴硕, 以庆, 弗敢, 今可令,
Nearest to 不能: 夹击, 指北, 五十匹, 谓关, 文安, 附贼, 夹口, 凛,
Nearest to 引兵: 借得, 感化, 计策, 退职, 家乐, 众军, 令永为, 赖汝众,
Nearest to 人马: 黄旗, 弓弩手, 欲待, 玄德便, 将星, 船只, 置库, 爵位,
Nearest to 关公: 玄德, 若进, 连夜, 盗, 粮草, 有水, 方圆, 抚,
Nearest to 一人: 奠仪, 魏续, 紧急, 英气, 动乎, 披银, 六十斤, 魏能,
Nearest to 玄德曰: 尽打, 犬子, 令二, 垂, 多名, 告立, 数口, 操执其,
Nearest to 袁绍: 天幸, 药, 新主, 劫迁, 表不信, 公之恨, 潘隐谓, 御车,
Nearest to 都督: 海若, 一语, 日应, 更择, 可筑, 荀有, 方有, 元叹,
Nearest to 却说: 拈阄, 选一舌辩, 法号, 布方, 投汉津, 世为, 水火之中, 部从止,
Nearest to 东吴: 引二子, 登程, 杜琼入, 其必出, 捉得, 忽觉, 乃亦布, 公惊,
Nearest to 刘备: 一万五千, 盛寒, 擅收, 当休, 之, 乃骤, 身骑, 特差,
Nearest to 赵云: 武骑, 非法, 劝表回, 严颜见, 老, 良图, 后面, 那人来,
Nearest to 二人: 承告, 四人助, 庙貌, 操欲, 植曰, 定睛, 新主, 河边,
可以看出，训练效果一般，很多判断近似词语都很牵强，所以我用了twitter的英文语料库训练了一次，英文语料库一共有1千万个单词，语料充足，同样训练了5轮，结果是这样的：
可以看出，语料充足的情况下word2vec学习得还是不错的。