Using pre-trained word embeddings

I used pre-trained word embeddings to initialize the weights of query and document embeddings like below code.

pre_trained_vectors = np.fromfile(path).reshape([-1, dim])
query_embedding_column = tf.feature_column.embedding_column(
    ...,
    initializer=tf.compat.v1.constant_initializer(pre_trained_vectors))

As a result, the training converged faster and the metric of result was pretty improved. I think using pre-trained vectors seems to be quite useful.

But the training time has increased.

The reason is that the weights are initialized with every evaluations. It takes unnecessarily long, because it will be restored from checkpoint file. And the data of pre-trained vectors is saved to meta file unnecessarily.

I’d like to know the way of initializing weights only at the first time of training.

[xuanhuiwang]

@muik, feel free to send a PR to add scaffold to _RankingHead.

[muik]

Thank you for the related link. I tried to use tf.train.Scaffold for initializing weights like the answer of the link. But scaffold argument could be passed to tf.estimator.EstimatorSpec, because the code where EstimatorSpec is instantiated is wraped by _RankingHead class on ranking/head.py.

So I had to modify _RankingHead class on head.py file to pass the scaffold argument. Now the problem is solved.

In addition, It would be nice to provide a way of passing arguments to creating EstimatorSpec.