With the help of awesome transformers library I have trained a multilabel classificator, which predicts topics for comments. I’m using TFDistilBertModel as a layer, wrapped in Keras functional API. So the architecture of a model is following: 2 input layers, then a layer, respawned as such:
distilbert_layer = TFDistilBertModel.from_pretrained(’/path/to/distilbert_model’), then some LSTM and Dense layers. The model is trained, everything is perfect at this time. Now comes the prediction time.
We serve our models on AWS containers, managed by kubernetes. Each time a user wants to generate topic of a new comment, we send a prediction task. To make a prediction I initiate model architecture, including respawning distilbert_layer as described above, then I load model weights, which I saved after training with the help of keras save_weights() function. Now the model is ready to predict.
The issue:
Each time a prediction task is sent, it consumes around 600mb of memory, which is not released after prediction. I assume there works a loadbalancer principle, so each time the task can be sent to a different process on the container, which uses different ram, which is why caching the model does not help. RAM quickly gets exhausted, then the worker just freezes and gets rebooted or initiates load of another worker.
Has anyone experienced such issues? Any help is very much appreciated.