Split neural network, load only needed part onto the GPU
I have a very, very big neural network and an Google Colab Pro Subscripting receiving my 16GB of GPU RAM. Unfortunately, this is not enough. My idea now is, to split the model (Unet) into the encoder and decoder part separately, and proceed like the following:

Load encoder to the GPU
Process the data through the encoder
Load encoder to the cpu, decoder to the GPU
Process the encoder output through the decoder
Load the decoder to the cpu aaaand repeat.

Is this in general possible? I coded an example but it wont work:
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=&quot;Epoch {}, train, loss={:.5f}&quot;.format(epoch, loss_value))
        running_loss.append(loss_value)

    return s.mean(running_loss)

For the setup / initialization I got the following:
DEVICE = &quot;cuda&quot;

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

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

Surely I understand the error, but I am sure that I push and pull everything at the right time onto and from the GPU... Or maybe there is a better way of splitting a model?

There are a couple of wrong things with this:

Your autocast block should include the forward on your model
You don't need to go back and forth from CPU to GPU and back, it's already a bottleneck with tensors, imagine with models.
Optional: what the heck is your UNet made of if you can't make it fit on a 16GB device?

I'd first try with standard solutions: reduce the batch size, include AMP (that I see you already included) or even DeepSpeed, which already does some CPU allocation, depending on the memory optimization level. Have a look at it, just in case, this may already solve your problems.
Answering to your question, the only feasible approach I see is to keep a model on cuda, the other on CPU and eventually move inputs/outputs.
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=&quot;Epoch {}, train, loss={:.5f}&quot;.format(epoch, loss_value))
        running_loss.append(loss_value)

    return s.mean(running_loss)

The main issue I see with this, is that your first model (residing on CPU) should not require the backpropagation part, otherwise it will probably break again. I doubt that PyTorch can build the backward graph between devices.
