memory leak in prediction tensorflow

TensorFlow
, ,
1

Recently I trained two MLP model and saved weights for future work.
Load model function contains this code to load models:

def creat_model_extractor(model_path, feature_count): try: tf.keras.backend.clear_session() node_list = [1024, 512, 256, 128, 64, 32] model = Sequential() model.add(Input(shape=(feature_count,))) for node in node_list: model.add(Dense(node, activation='relu')) model.add(Dropout(0.2)) model.add(LayerNormalization()) model.add(Dense(16, activation='relu')) model.add(LayerNormalization()) model.add(Dense(1, activation='sigmoid')) model.load_weights(model_path) model.trainable = False for layer in model.layers: layer.trainable = False except Exception as error: logger.warning(error, exc_info=True) return None return model

And this my prediction code

@tf.function def inference(model, inputs): return tf.stop_gradient(model(inputs, training=False)) predictions = inference(setting.SMALL_MODEL, small_blocks_normal) small_blocks['label'] = (predictions > 0.5).numpy().astype(int)

I have question about a memory leak.

After predicting labels, the memory is never released—even after the function finishes executing.

What should I do? I tried the following code, but it didn’t work.

I use CPU and don’t have any gpu.

import gc tf.keras.backend.clear_session() del predictions gc.collect()
3

Hi @Leaorad_LapTop , If K.clear_session() didn’t resolve the problem, try adding garbage collection after each call with _ = gc.collect(). This should help to remain the memory usage constant. Thank you!