I'm developing a test module to see how many FastAI models I can load before my memory crashes.
At the moment, I'm having an issue removing the models from memory once it's full. This is what I've attempted:
import gc
from fastai.vision.all import load_learner
class ClassificationInference:
def __init__(self):
self.model_list = []
def load_model(self, model_name: str) -> None:
model_file = os.path.join("trained_model",
"saved_model", f"{model_name}.pkl")
model = load_learner(model_file)
self.model_list.append(model)
def remove_all_models(self) -> None:
for i in range(len(self.model_list)):
if i >= len(self.model_list):
break
del self.model_list[i]
del self.model_list
self.model_list = None
torch.cuda.empty_cache()
gc.collect()
self.model_list = []
I'm measuring the RAM usage with the following commands:
max_gpu = torch.cuda.get_device_properties(None).total_memory
max_ram = psutil.virtual_memory().total
current_gpu = torch.cuda.memory_allocated(None)
current_ram = psutil.virtual_memory().used
percentage_gpu = (current_gpu / max_gpu) * 100
percentage_ram = (current_ram / max_ram) * 100
print(f"Percentage before start: (RAM: {percentage_ram}, "
f"GPU: {percentage_gpu})")
After I execute remove_all_models
, the RAM percentage is the same as before.
How can I delete the models from memory? Thanks in advance.
A quote from "Expert Python Programming":
So you're on the right track, but there is also one more thing you need to be aware of:
You would need to do a del on all other objects that interact with your models, since those are also holding a reference to the same objects. Calling del will only decrement the reference count and remove the particular reference from usage, but the actual in memory object is not garbage collected until the reference count hits 0.