我有一个神经网络，它以(positive_document, negative_document)的形式接受输入。输入必须来自TFRecordDataset，因为我希望在TPU上训练模型。我正在使用成对的铰链损耗。
问题：
我最终将重复数据存储在TFRecord中。例如，存储在磁盘上的TFExample是：
(posdoc_1, negdoc_1),
(posoc_1, negdoc_2),
.
.
(posodoc_1, negdoc_n)
(posdoc_2, negdoc_1),
(posdoc_2, negdoc_2),
.
posdoc_1和posdoc_2不必要地重复。
解决方案的想法
如果我可以将TFExample存储为：
(posdoc_1_id, negdoc_1_id)
(posdoc_1_id, negdoc_2_id)
.
.
存储ID会非常便宜，希望我可以在培训期间将其转换为实际的doc表示形式。像这样：
for pos_id, neg_id in id_tfrecords:
posdoc = data_tfrecords[pos_id]
negdoc = data_tfrecords[neg_id]
id_tfrecords包含(posid, negid)对，而data_tfrecords包含实际的文档表示形式
由于多种原因，上述内容显然行不通（TFRecordDataset不能下标），但是现在您知道我的想法了。
我已经尝试过的内容
对于适合内存的数据集，我在磁盘上写入了(posid, negid)对，并在内存中保持了一个张量，该张量可以使我得到文档表示。这或多或少像keras / pytorch中的EmbeddingLayer。例如，这就是我将训练三元组转换为实际表示形式的方式：
for posid, negid in training_data:
pos_doc = docid_to_doc[posid] # docid_to_doc is a tensor. Assume docid is an index into this tensor
.
.
我很快遇到了问题：
显然，张量在内存中不能超过2GB。
我尝试使用通过自定义映射初始化的tf.Keras.layers.Embedding，但是仅适用于将idx转换为向量。我希望将idx转换为2D矩阵。
最终目标是：
任何想法都将不胜感激。预先感谢。