[Python] Ram Memory leak when scripting a Sampling Trainer for a Bert Encoder and LSTM Decoder...

I wrote the module attached below. However, I notice a constant increase of RAM until I get an out of memory error. The code runs on CPU without a problem (except the slow training time). It can finish the first training step, however the memory doesn't get release after the first training cycle and it crashes out. I would be extremely grateful if someone can solve the problem or guess what the problem is.

import tensorflow as tf
import gc
from config import SAMPLING_PROBABILITY
#tf.config.run_functions_eagerly(True)

class ScheduledSamplingTrainer(tf.keras.Model):
def __init__(self, model):
super().__init__()
self.model = model
self.schedule_prob = tf.Variable(SAMPLING_PROBABILITY, trainable=False, dtype=tf.float32)
self.loss_fn = tf.keras.losses.MeanSquaredError()
self.metric = tf.keras.metrics.MeanAbsoluteError()

#@tf.function
def train_step(self, data):
x, y_true = data
(encoder_input_ids, encoder_attention_mask, decoder_inputs) = x
y_true = tf.cast(y_true, tf.float32)
batch_size = tf.shape(encoder_input_ids)[0]
pose_dim = tf.shape(y_true)[-1]
seq_len = int(decoder_inputs.shape[1])

# mask per batch: (batch, seq_len, 1)
mask = tf.cast(tf.expand_dims(tf.reduce_any(y_true != 0, axis=-1), -1), tf.float32)

go_frame = tf.fill((batch_size, 1, pose_dim), tf.constant(0, dtype=tf.float32))
outputs_ta = tf.TensorArray(dtype=tf.float32, size=seq_len, dynamic_size=False)

use_pred_mask = tf.random.uniform((seq_len,), dtype=tf.float32) < self.schedule_prob

with tf.GradientTape() as tape:
encoder_outputs = self.model.bert_encoder(
input_ids=encoder_input_ids,
attention_mask=encoder_attention_mask,
training=True
)

mean_embedding = tf.reduce_mean(encoder_outputs.last_hidden_state, axis=1)
mean_embedding = tf.cast(mean_embedding, tf.float32)
initial_h = self.model.decoder_initial_state_h(mean_embedding)
initial_c = self.model.decoder_initial_state_c(mean_embedding)
encoder_states = [initial_h, initial_c]

decoder_input = go_frame
for t in range(seq_len):


decoder_outputs, h, c = self.model.decoder_lstm(
self.model.masking(decoder_input),
initial_state=encoder_states
)
encoder_states = [h,c]

next_frame = self.model.decoder_dense(decoder_outputs[:, -1:, :])
next_frame = tf.cast(next_frame, outputs_ta.dtype)
# scheduled sampling

next_input = tf.where(use_pred_mask[t], next_frame, y_true[:, t:t+1, :])

decoder_input = next_input

outputs_ta = outputs_ta.write(t, next_frame)

outputs = outputs_ta.stack() # [seq_len, batch, 1, pose_dim]
outputs = tf.transpose(outputs, [1, 0, 2, 3]) # [batch, seq_len, 1, pose_dim]
outputs = tf.squeeze(outputs, axis=2)

masked_outputs = outputs * mask
masked_y_true = y_true * mask


# --- PER-TIMESTEP LOSS ---
loss = 0.0
for t in range(seq_len):
loss += tf.reduce_mean(tf.square(masked_y_true[:, t, :] - masked_outputs[:, t, :]))
loss /= tf.cast(seq_len, tf.float32)

#loss = tf.cast(loss, self.model.trainable_variables[0].dtype)
grads = tape.gradient(loss, self.model.trainable_variables)
self.optimizer.apply_gradients(zip(grads, self.model.trainable_variables))

# masked metric update
self.metric.update_state(masked_y_true, masked_outputs)
del outputs_ta, decoder_input, decoder_outputs, next_frame, masked_outputs, masked_y_true, tape
gc.collect()
return {"loss": loss, "mae": self.metric.result()}

Continue reading...