Memory Leak / Утечка памяти

[Freeeeeet]

Столкнулся с проблемой утечки памяти на модели tts. При каждом обращении к модели занятая оперативка растет на ~100МБ.

Вызывается сразу несколько инстансов, работает через FastAPI, но я думаю, можно обойтись без деталей и оставил только то, что связано с моделью(если что-то ещё нужно будет докинуть - докину)
вызываю примерно вот таким образом, чтобы получить итоговое аудио:

model = get_model()
    file_path = generate_audio_file(model, query.id, query.text_to_speech)
    return_model(model)

А вот тут в различных файликах описано, что я, собственно, вызываю:

device = 'cpu'
model_pool = ModelPool(size=2, device=device)



def get_model():
    return model_pool.get_model()


def return_model(model):
    model_pool.return_model(model)

class ModelPool:
    def __init__(self, size, device='cpu'):
        self.pool = queue.Queue()
        self.device = torch.device(device)  # Устройство, на котором будет выполняться модель
        self.load_models(size)

    def load_models(self, size):
        for _ in range(size):
            model, __ = torch.hub.load(repo_or_dir='snakers4/silero-models',
                                   model='silero_tts',
                                   language='ru',
                                   speaker='v3_1_ru')
            model.to(self.device)
            self.pool.put(model)

    def get_model(self):
        return self.pool.get()

    def return_model(self, model):
        self.pool.put(model)

def generate_audio_file(model, id, text: str, speaker='xenia', sample_rate=24000) -> str:
    file_path = f"some_filename--{id}.wav"
    model.save_wav(text=text, speaker=speaker, sample_rate=sample_rate, audio_path=file_path)
    return file_path

Не сталкивались с таким?
torch.no_grad() пробовал добавлять в различные места кода, не помогло

[dumonten]

Привет! У меня было точно так же, когда я использовала v2, RAM выедалась до 1 ГБ и выше. Не понимала в чем проблема, использовала и torch._C._jit_set_profiling_mode(False), и torch.set_grad_enabled(False), вызывала и gc.collect(), и torch.cuda.empty_cache() - ничего не помогало. С этим, видимо, все сталкиваются, так именно моделька написана - и, как я понимаю, это не исправить. Поставила v4 - и все стало ок

[MrYNikita]

@dumonten, приветствую! Ознакомился с Вашим кодом, для Вас проблема была решена и утечка памяти устранена?

[Freeeeeet]

Нет же, все описал, как выглядят примеры моего кода
Все это крутится в докере и дёргается через фастапи
И каждый раз, когда дёргается модель таким образом, как описано выше - начинает течь память.
Принципиально использовать версию модели v3.1

[dumonten]

Нет же, все описал, как выглядят примеры моего кода Все это крутится в докере и дёргается через фастапи И каждый раз, когда дёргается модель таким образом, как описано выше - начинает течь память. Принципиально использовать версию модели v3.1

Может попробуй еще позапускать код с apply_tts, а не через save_wav? Вот пример:

audio = model.apply_tts(ssml_text=ssml_sample,
speaker=speaker,
sample_rate=sample_rate)

Только подчищай тензор, который он генерит после сохранения.

Я лично, v3 не использовала, только v2 и v4. Поэтому и порекомендовала то решение, где у меня все исправилось.

[VitalyPetrov]

Можно использовать io.BytesIO объект для сохранения wav в него и перезаписывать его при каждом обращении к tts

import io

bts = io.BytesIO()

model.save_wav(
    ssml_text=ssml,
    sample_rate=48000
    speaker='<some-speaker-name>',
    put_accent=True,
    put_yo=True,
    audio_path=bts,
)

Соответственно, при каждом вызове метода .save_wav() bts будет ссылаться на новую область памяти, а старую сборщик мусора будет удалять.