How to delete a list of FastAI models from memory?

57 views Asked by At

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.

1

There are 1 answers

1
mighty_mike On

A quote from "Expert Python Programming":

The approach of such a memory manager is roughly based on a simple statement: If a given object is not referenced anymore, it is removed. In other words, all local references in a function are removed after the interpreter:

• Leaves the function

• Makes sure the object is not being used anymore.

Under normal conditions, the collector will do a nice job. But a del call can be used to help the garbage collector by manually removing the references to an object manually.

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.