How to unload or delete a tf.saved_model.load(path)

36 Views Asked by At

Using API I am loading a tf saved_model each time for processing the audio. AS each time the model loaded for processing the audio leads me to the memory leakage. So Is there any way to unload the model from the memory or delete it from the memory. the del and followed by gc.collect function does not work for me. Here is the code: #!/usr/bin/env python3

-- coding: utf-8 --

import os os.environ["CUDA_VISIBLE_DEVICES"]="-1" import soundfile as sf import gc import numpy as np import tensorflow as tf import librosa from numba import cuda

def noise_supression(dtln_model_path,audio_data,sf):

block_len = 512
block_shift = 128
model = tf.saved_model.load(dtln_model_path)#'/kaldi_docker/echo/Audio_Analytics/DTLN/DTLN/pretrained_model/dtln_saved_model')
infer = model.signatures["serving_default"]
audio = audio_data
fs = sf 
if fs != 16000:
    raise ValueError('This model only supports 16k sampling rate.')
# preallocate output audio
out_file = np.zeros((len(audio)))
# create buffer
in_buffer = np.zeros((block_len))
out_buffer = np.zeros((block_len))
# calculate number of blocks
num_blocks = (audio.shape[0] - (block_len-block_shift)) // block_shift
# iterate over the number of blcoks        
for idx in range(num_blocks):
    #print("here")
    # shift values and write to buffer
    in_buffer[:-block_shift] = in_buffer[block_shift:]
    in_buffer[-block_shift:] = audio[idx*block_shift:(idx*block_shift)+block_shift]
    # create a batch dimension of one
    in_block = np.expand_dims(in_buffer, axis=0).astype('float32')
    # process one block
    print(type(infer))
    out_block= model_infer(tf.constant(in_block))['conv1d_1']
    # shift values and write to buffer
    out_buffer[:-block_shift] = out_buffer[block_shift:]
    out_buffer[-block_shift:] = np.zeros((block_shift))
    out_buffer  += np.squeeze(out_block)
    # write block to output file
    out_file[idx*block_shift:(idx*block_shift)+block_shift] = out_buffer[:block_shift]
        
print('Noise Supression Processing finished.')
del(in_buffer)      #mmk 30May23    Release memory allocated
del(out_buffer)     #mmk 30May23    Release memory allocated
del model
gc.collect()
return out_file

I want to load the model once and use it multiple time. how this could be possible.

0

There are 0 best solutions below