拆分神经网络，只将需要的部分加载到 GPU 上

Question

我有一个非常非常大的神经网络和一个接收我的 16GB GPU RAM 的 Google Colab Pro 订阅。不幸的是，这还不够。我现在的想法是，将模型（Unet）分别拆分为编码器和解码器部分，如下进行：

1. 将编码器加载到 GPU
2. 通过编码器处理数据
3. 将编码器加载到 CPU，将解码器加载到 GPU
4. 通过解码器处理编码器输出
5. 将解码器加载到 cpu aaa 并重复。

这通常是可能的吗？我编写了一个示例，但它不起作用：

def train(epoch, loader, loss_fn, optimizer, scaler, model1, model2):
    model1.train()
    model2.train()

    loop = prog(loader)

    running_loss = []

    for batch_index, (data, target) in enumerate(loop):
        optimizer.zero_grad(set_to_none=True)

        model1 = model1.to(DEVICE)

        data, skip_connections = model1(data.to(DEVICE))
        model1 = model1.cpu()

        model2 = model2.to(DEVICE)
        data = model2(data, skip_connections)

        model2 = model2.cpu()

        target = target.to(DEVICE)

        with torch.cuda.amp.autocast():
            loss = loss_fn(data, target)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        model1 = model1.to(DEVICE)

        loss_value = loss.item()

        loop.set_postfix(info="Epoch {}, train, loss={:.5f}".format(epoch, loss_value))
        running_loss.append(loss_value)

    return s.mean(running_loss)

对于设置/初始化，我得到了以下信息：

DEVICE = "cuda"

model1 = UNET_FIRST_HALVE(in_channels=4).to(DEVICE)
model2 = UNET_SECOND_HALVE(out_channels=NUM_CLASSES).cpu()

for epoch in range(epochs_done + 1, num_epochs + 1):
    training_loss = train(..., model1, model2)
    .
    .
    .

我收到以下错误：

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

我当然理解错误，但我确信我会在正确的时间将所有内容推入和拉出 GPU...或者也许有更好的分割模型的方法？

Answer 1

这里有几个错误的地方：

1. 你的 autocast 块应该在你的模型中包含 forward
2. 您不需要在 CPU 和 GPU 之间来回切换，这已经是张量的瓶颈，想象一下模型。
3. 可选：如果您的 UNet 无法安装在 16GB 设备上，那么它到底是由什么组成的？

我首先尝试使用标准解决方案：减少批量大小，包括 AMP（我看到您已经包括在内）甚至 DeepSpeed，它已经根据内存优化级别进行了一些 CPU 分配。 Have a look at it，以防万一，这可能已经解决了你的问题。

回答您的问题，我认为唯一可行的方法是将模型保留在 cuda 上，将另一个模型保留在 CPU 上并最终移动输入/输出。

def train(epoch, loader, loss_fn, optimizer, scaler, model1, model2):
    model1.to(device1)
    model2.to(device2)
    model1.train()
    model2.train()

    loop = prog(loader)

    running_loss = []

    for batch_index, (data, target) in enumerate(loop):
        optimizer.zero_grad(set_to_none=True)
        # now data nd skip are on device1
        data, skip_connections = model1(data.to(device1))
        # we need to move to device2
        data = data.to(device2)
        skip = skio_connections.to(device2)
        data = model2(data, skip)
        # everything should be on 2
        # move the target to the same device
        target = target.to(device2)
        loss = loss_fn(data, target)
        # backprop
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        loss_value = loss.item()
        loop.set_postfix(info="Epoch {}, train, loss={:.5f}".format(epoch, loss_value))
        running_loss.append(loss_value)

    return s.mean(running_loss)

我看到的主要问题是，您的第一个模型（驻留在 CPU 上）不应该需要反向传播部分，否则它可能会再次中断。我怀疑 PyTorch 能否在设备之间构建后向图。