pytorch - 如何检测pytorch中的内存泄漏？

Question

我正在从头开始训练 VGG11，尽管我完全清理了 GPU，但我的 GPU 中一直出现“内存不足”的情况。

    self.criterion = nn.CrossEntropyLoss()
    lr_dict = {100: 0.01, 150: 0.001}
    self.optimizer = optim.SGD(self.model.parameters(), lr=0.1, momentum=0.9)
    result_per_epoc = list()
    train_acc = 0.0
    train_loss = 0.0
    for epoch in range(epochs):
        debug_memory(f"Train epoc {epoch} START", clear_mem=True)

        if epoch in lr_dict.keys():
            for gopt in self.optimizer.param_groups:
                gopt['lr'] = lr_dict[epoch]
        epoch_time_start = datetime.now()

        for batch_idx, (images, labels) in tqdm(enumerate(iter(traindata_loader)), total=len(traindata_loader)):
            images = images.to(self.device)
            labels = labels.to(self.device)
            self.optimizer.zero_grad()
            try:
                outputs = self.model(images)
            except Exception as e:
                debug_memory(f"On crash", clear_mem=True)
                print(f"Batch ID: {batch_idx}")
                raise e
            loss = self.criterion(outputs, labels)
            loss.backward()
            self.optimizer.step()

            train_accuracy = accuracy(outputs, labels)
            train_acc += float(train_accuracy)
            train_loss += float(loss)

            del images, labels, loss, outputs
            clear_memory()
        debug_memory(f"Train epoc {epoch} TRAIN", clear_mem=True)
        train_acc /= (batch_idx + 1)
        train_loss /= (batch_idx + 1)

直到行 "debug_memory(f"Train epoc {epoch} START", clear_mem=True)" 我使用了 3,499/11,019 (3​​1%) 的内存。但是经过几批之后，我的代码在内存上崩溃了，即使我删除了我添加到 GPU 的所有内容，并且在“clear_memory”中我这样做了：

torch.cuda.empty_cache()
gc.collect()
with torch.no_grad():
    torch.cuda.empty_cache()
gc.collect()

我使用的模型是来自 torchvision.models 的“vgg11_bn”

我还缺少什么？

标签： pytorch