由于原生tf缺少对大规模稀疏embedding的支持,为了支持大规模id特征,在项目里使用了Dynamic Embedding,但是在尝试使用分布式进行训练时却产生了如下报错:
Waiting for model to be ready. Ready_for_local_init_op: Variables not initialized: beta1_power, beta2_power, ready: None
网上有很多类似的报错但是解决方案都不适合这个场景,在浏览了很多的搜索结果后突然发现https://github.com/tensorflow/recommenders-addons/issues/38这个issue,才解决了问题。
问题原因
问题起源于dynamic embedding的embedding_lookup。
beta1_power和beta2_power是adam创建的两个参数。adam优化器在创建这两个参数时会根据var_lisr名字排序,选择一个最靠前的变量,按照这个变量的设备创建beta1_power和beta2_power。这样做的目的是让所有worker都能在相同的ps上创建这两个参数[1]。
TrainableWrapper是在embedding_lookup中创建的一个ResourceVariable,用于辅助传递embedding和计算梯度,位于worker上。如果某个TrainableWrapper正好位于var_list的第一个,那么adam就会按照TrainableWrapper的设备来创建beta1和beta2这两个参数,也就是worker,导致出错。
参考:
- [1] adam.py