PyTorch GPU内存不足

Question

我正在PyTorch中运行一个评估脚本。我有许多经过训练的模型（*.pt文件），我将其加载并移动到GPU，总共占用270MB的GPU内存。我使用的批量大小为1。对于每个示例，我加载一个图像并将其移动到GPU。然后，根据样本，我需要运行一系列经过训练的模型。有些模型以张量作为输入和输出。其他模型的输入是张量，输出是字符串。序列中的最终模型总是有一个字符串作为输出。中间张量临时存储在字典中。当模型使用张量输入时，使用del将其删除。不过，我注意到，每次采样后，GPU内存都会不断增加，直到整个内存都满了为止。

下面是一些伪代码，让您更好地了解正在发生的事情：

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 ...

我不知道为什么GPU内存不足。有什么想法吗？

Answer 1

尝试添加火炬。库达。删除后清空_缓存（）