小型 PyTorch 网络需要近 3 GB 的 RAM 在 MNIST 上进行训练

Question

我在使用 PyTorch 时遇到问题。 我必须在 PyTorch 自定义模型上运行一些实验，考虑到我必须训练很多模型，我尝试使用

ray[tune]

并行运行它们。我已经编写了代码，并且运行得“很好”，但问题是我无法并行化太多，因为系统内存不足。我发现训练中的每个 PyTorch 模型都使用类似 2/2.5 GB 的 RAM。问题是我正在运行非常小的网络，总共 3 层，并且我正在 MNIST 上运行测试，所以一切都应该是非常轻量级的，为什么训练这样小的网络需要使用这么多 RAM？ 这个当我尝试并行运行代码时，这当然成为一个真正的问题，因为仅仅 4 个并行训练最终就会占用超过 8 GB 的 RAM。 这正常吗？

我当然认为我的代码有问题，所以我在 PyTorch 中创建了一个独立的脚本，并在 MNIST 上创建了一个小型网络，以测试它最终使用了多少内存；这是代码，它是真正标准的设计：

# FILENAME: selfcontained_pytorch_tester.py
# Importing the necessary modules for the training and testing of a simple network in pytorch with MNIST

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

# Checking if the GPU is available and setting the device
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"Using {device} device")

# Loading the data
training_data = datasets.MNIST(
    root="Data_SelfContained_Pytorch_Testing/data_train_MNIST",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.MNIST(
    root="Data_SelfContained_Pytorch_Testing/data_test_MNIST",
    train=False,
    download=True,
    transform=ToTensor()
)


# Creating the dataloaders
dataloader_train = DataLoader(training_data, batch_size=64)
dataloader_test = DataLoader(test_data, batch_size=64)


# Defining the network
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
            nn.ReLU()
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits


# Creating the model
model = SimpleNN().to(device)

# Defining the loss function
loss_fn = nn.CrossEntropyLoss()

# Defining the optimizer Adam
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)


# Defining the training function
def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # Compute prediction error
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


# Defining the testing function
def test(dataloader, model, loss_fn):
    size = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= size
    correct /= size
    print(f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")


# Training and testing the model
epochs = 5
for t in range(epochs):
    print(f"Epoch {t + 1}\n-------------------------------")
    train(dataloader_train, model, loss_fn, optimizer)
    test(dataloader_test, model, loss_fn)
print("Done!")

然后我运行脚本，只需使用

htop

监控 RAM 使用情况，即使是这个标准的、非常简单的 pytorch 脚本最终也会使用大约 2.5 GB 的 ram（甚至更多！）。

这是

htop

的屏幕：

如您所见，我们有巨大的

VIRT

并且

RES

约为。 2.9 GB。

发生什么事了？这对于 PyTorch 来说正常吗？

我从在 TensorFlow 上运行网络的朋友那里听说，对于训练这么小的网络来说，这些内存要求是荒谬的。这是真的吗？

---

更新：我修改了代码，看看是否应该归咎于

cuda

设备（因此GPU），事实上，如果我使用

device = 'cpu'

运行代码，脚本使用的RAM量会下降到大约 500MB；因此，使用 Nvidia GPU (cuda) 时，实质上会使用超过 2 GB 的额外 RAM 内存。又问：这是怎么回事？这是正常的吗？我正在使用我安装的 pytorch 版本：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

nvidia-smi

的输出是：

所以有两个不同的cuda版本，但我读到它们应该是兼容的。现在我当然想知道这是否是问题所在，但我不这么认为，我在其他配置下运行代码获得了相同的结果..

有问题吗？使用 cuda 时这种 RAM 分配是必要的和预期的行为吗？如果是这样，当使用

ray

同时运行多个训练时，有没有办法“支付此 RAM 价格”一次？

Answer 1

我想我已经找到解决方案了！

在所有相关调用

torch

之前，您可以放置

# At the start of the script
import os
os.environ['CUDA_MODULE_LOADING'] = 'LAZY'

这应该告诉

torch

不要在 RAM 中加载任何不必要的组件。 但是，这里有一个问题，只有当你有 cuda 版本 11.8 或更高版本时，这才有效。我必须更新我的驱动程序才能使其正常工作。执行此操作时，我使用

cuda

的内存需求从原来的 2.5 GB 下降到 700 MB。

就是这样！我认为这就是解决方案；现在一切正常。