在張量流圖中加載大的numpy矩陣作爲分區變量

Question

想象一下，我有大的預訓練嵌入，我可以加載爲numpy數組，例如形狀爲[3000000,200]。這個矩陣的大小是大於2GB，所以使用此代碼：

data = np.zeros(shape=(3000000, 200)) 
variable = tf.get_variable(
    "weigths", 
    [3000000, 200], 
    initializer=tf.constant_initializer(data)) 

session = tf.Session() 
session.run(tf.global_variables_initializer()) 

我有錯誤ValueError: Cannot create a tensor proto whose content is larger than 2GB.

我可以tf.assign和佔位符加載它，但由於某些原因，我想使用這個嵌入權重的分區版本。分配和佔位符的方式已關閉，因爲分區變量不適用於分配操作：NotImplementedError: assign() has not been implemented for PartitionedVariable.。

是否有可能做這樣的事情？

Answer 1

SOLUTION

這是醜陋的，但它的工作原理：

def init_partitioned(session, var_name, data): 
    partitioned_var = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=var_name + "/part_\d+:0") 
    print("For {} founded {} parts".format(var_name, len(partitioned_var))) 

    dtype = partitioned_var[0].dtype 
    part_shape = partitioned_var[0].get_shape().as_list() 
    part_shape[0] = None 

    init = tf.placeholder(dtype, part_shape) 
    offset = 0 
    for idx, part in enumerate(partitioned_var): 
     init_op = tf.assign(part, init) 
     numRowsInPart = int(part.get_shape()[0]) 
     session.run(init_op, feed_dict={init: data[offset:offset + numRowsInPart]}) 
     offset += numRowsInPart 

Answer 2

嘗試：

import numpy as np 
import tensorflow as tf 

data = np.zeros(shape=(3000000, 200)) 

ph = tf.placeholder(tf.float32, shape=(3000000, 200)) 
variable = tf.Variable(ph) 

session = tf.Session() 
session.run(tf.global_variables_initializer(), feed_dict={ph:data})