PyTorch GPU 内存不足

Question

I am running an evaluation script in PyTorch. I have a number of trained models (*.pt files), which I load and move to the GPU, taking in total 270MB of GPU memory. I am using a batch size of 1. For every sample, I load a single image and also move it to the GPU. Then, depending on the sample, I need to run a sequence of these trained models. Some models have a tensor as input and as output. Other models have a tensor as input, but a string as output. The final model in a sequence always has a string as output. The intermediary tensors are temporarily stored in a dictionary. When a model has consumed a tensor input, it is deleted using del . Still, I notice that after every sample, the GPU memory keeps increasing until the entire memory is full.

Below is some pseudocode to give you a better idea of what is going on:

with torch.no_grad():
    trained_models = load_models_from_pt() # Loaded and moved to GPU, taking 270MB
    model = Model(trained_models) # Keeps the trained_models in a dictionary by name
    for sample in data_loader:
        # A sample contains a single image and is moved to the GPU
        # A sample also has some other information, but no other tensors
        model.forward(sample)

class Model(nn.Module)
    def __init__(self, trained_models):
        self.trained_models = trained_models
        self.intermediary = {}

    def forward(sample):
        for i, elem in enumerate(sample['sequence']):
             name = elem['name']
             in = elem['input']
             if name == 'a':
                model = self.trained_models['a']
                out = model(self.intermediary[in])
                del self.intermediary[in]
                self.intermediary[i] = out
             elif name == 'b':
                model self.trained_models['b']
                out = model(self.intermediary[in])
                del self.intermediary[in]
                self.intermediary[i] = out
             elif ...

I have no idea why the GPU is out of memory. Any ideas?

Answer 1

尝试在 del 之后添加 torch.cuda.empty_cache()