Исправление ResourceExhaustedError, которого не было на TF1.1.4, но на 2.0.0

Question

Я пробую Tensorflow 2.0.0 и хочу посмотреть, как он работает в моем текущем проекте.

В Tensorflow 1.14 у меня не было проблем с запуском этой модели с ее размером и таким же оборудованием.

def __init__(self, num_input=48959, num_hid1=2048, num_hid2=512, latent_rep=196 , lr=0.00005, dtypetf=tf.float32, dtypenp=np.float32):

        self.graph = tf.Graph()
        actf=tf.nn.relu
        with self.graph.as_default():


            self.X=tf.placeholder(shape=[None,num_input] ,dtype=dtypetf )
            initializer=tf.variance_scaling_initializer(dtype=dtypetf)

            w1=tf.Variable(initializer([num_input,num_hid1]),dtype=dtypetf)
            w2=tf.Variable(initializer([num_hid1,num_hid2]),dtype=dtypetf)
            w3=tf.Variable(initializer([num_hid2,latent_rep]),dtype=dtypetf)
            w4=tf.Variable(initializer([latent_rep,num_hid2]),dtype=dtypetf)
            w5=tf.Variable(initializer([num_hid2,num_hid1]),dtype=dtypetf)
            w6=tf.Variable(initializer([num_hid1,num_input]),dtype=dtypetf)

            b1=tf.Variable(tf.zeros(num_hid1,dtype=dtypenp),dtype=dtypetf)
            b2=tf.Variable(tf.zeros(num_hid2,dtype=dtypenp),dtype=dtypetf)
            b3=tf.Variable(tf.zeros(latent_rep,dtype=dtypenp),dtype=dtypetf)
            b4=tf.Variable(tf.zeros(num_hid2,dtype=dtypenp),dtype=dtypetf)
            b5=tf.Variable(tf.zeros(num_hid1,dtype=dtypenp),dtype=dtypetf)
            b6=tf.Variable(tf.zeros(num_input,dtype=dtypenp),dtype=dtypetf)

            hid_layer1=actf(tf.matmul(self.X,w1)+b1)
            hid_layer2=actf(tf.matmul(hid_layer1,w2)+b2)
            hid_layer3=actf(tf.matmul(hid_layer2,w3)+b3)
            hid_layer4=actf(tf.matmul(hid_layer3,w4)+b4)
            hid_layer5=actf(tf.matmul(hid_layer4,w5)+b5)
            self.output=actf(tf.matmul(hid_layer5,w6)+b6)


def opt(self, loc='float32_mem.npy' , epoches=250, batch_size=128, num_batches=100):

        imported = np.load(loc)
        with tf.Session(graph = self.graph) as sess:
            sess.run(self.init)
            for epoch in range(epoches):
                for iteration in range(num_batches):
                    sector = (iteration*batch_size)%1024
                    sector_end = (sector+batch_size-1)%1024
                    X_batch = imported[sector:sector_end]
                    sess.run(self.train,feed_dict={self.X:X_batch})

Код Tensorflow 2.0.0 Keras (не работает)

keras.backend.set_floatx('float32')
model = keras.Sequential([
    keras.layers.Dense(48959,activation='relu'),
    keras.layers.Dense(1024, activation='relu'),
    keras.layers.Dense(512, activation='relu'),
    keras.layers.Dense(196, activation='relu'),
    keras.layers.Dense(512, activation='relu'),
    keras.layers.Dense(1024, activation='relu'),
    keras.layers.Dense(48959)
])
model.compile(optimizer='adam',loss='mean_squared_error',metrics=['accuracy'])

model.fit(x=training_data, y=training_data, epochs=250,batch_size=128 )

или что я уже пробовал:

epochs=250
batch_size=128
num_batches=100

imported  = np.load('float32_mem_3.npy')
for epoch in range(epochs):
    for iteration in range(num_batches):
        sector = (iteration*batch_size)%1024
        sector_end = (sector+batch_size-1)%1024
        X_batch = imported[sector:sector_end]
        model.train_on_batch(x=X_batch, y=X_batch)

Ошибка, которую я получаю: ResourceExhaustedError: OOM when allocating tensor with shape[48959,48959] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc [Op:RandomUniform]

Полагаю, это не актуально, но я использую 980 ti с 6 ГБ видеопамяти.