Есть ли способ частичной загрузки весов в память в Tensorflow или Pytorch, в тех случаях, когда тренируется только часть весов за шаг?

Question

Мне интересно, что есть способ сохранить большинство весов только в памяти диска и загружать их в память графического процессора только тогда, когда они собираются или собираются быть обученными.

Я обучаю встраивания для представления элементов (например, word2vec, Glove и т. Д.), И у меня есть более 40 миллионов элементов с размерами встраивания, идеально равными 256. На каждом этапе обучения обучается / обновляется только небольшая часть параметров. .

Моя среда Tensorflow аварийно завершает работу при инициализации входной матрицы внедрения и матрицы внедрения softmax.

Если полезно, вот как выглядит мой график для справки:

batch_size = 16384
embedding_size = 256
num_inputs =5
vocabulary_size = 40000000

num_sampled = 64 # Number of negative examples to sample.

graph = tf.Graph()

with graph.as_default() , tf.device('/gpu:0'):

    train_dataset = tf.placeholder(tf.int32, shape=[batch_size, num_inputs ])
    train_labels = tf.placeholder(tf.int32, shape=[batch_size, 1])

    embeddings = tf.get_variable( 'embeddings', dtype=tf.float32,
        initializer= tf.random_uniform([vocabulary_size, embedding_size], -1.0, 1.0, dtype=tf.float32) )

    softmax_weights = tf.get_variable( 'softmax_weights', dtype=tf.float32,
        initializer= tf.truncated_normal([vocabulary_size, embedding_size],
                             stddev=1.0 / math.sqrt(embedding_size), dtype=tf.float32 ) )

    softmax_biases = tf.get_variable('softmax_biases', dtype=tf.float32,
        initializer= tf.zeros([vocabulary_size], dtype=tf.float32),  trainable=False )

    embed = tf.nn.embedding_lookup(embeddings, train_dataset) #train data set is

    embed_reshaped = tf.reshape( embed, [batch_size*num_inputs, embedding_size] )

    segments= np.arange(batch_size).repeat(num_inputs)

    averaged_embeds = tf.segment_mean(embed_reshaped, segments, name=None)

    loss = tf.reduce_mean(
        tf.nn.sampled_softmax_loss(weights=softmax_weights, biases=softmax_biases, inputs=averaged_embeds, 
                                   sampled_values=tf.nn.uniform_candidate_sampler(true_classes=tf.cast(train_labels, tf.int64), num_sampled=num_sampled, num_true=1, unique=True, range_max=vocabulary_size, seed=None),
                                   labels=train_labels, num_sampled=num_sampled, num_classes=vocabulary_size)) 

    optimizer = tf.train.AdagradOptimizer(1.0).minimize(loss) 

Моя среда Tensorflow может обрабатывать только ~ 5 миллионов предметов. Все вложения инициализированы и находятся в памяти и готовы к обучению. Тем не менее, для каждого обучающего шага единственными вложениями, которые обучаются, являются: вложения softmax для каждой метки в пакете, вложения ввода для всех входов в пакете и вложение softmax для отрицательных образцов для всех меток в партия

Таким образом, на каждом этапе обучения тренируется только очень небольшая часть от общего веса.

Мне интересно, что есть способ сохранить большинство весов только в памяти диска и загружать их в память графического процессора только тогда, когда они собираются или собираются быть обученными.

Пока что я не видел ничего подобного в Tensorflow. Можно ли это сделать в Pytorch?