Moving PyTorch backend tensors between GPU and CPU

[Dxyk]

Hello,

I'm working with PhiFlow's PyTorch backend. At the beginning of my code, I specify the GPU as my default device by the following snippet.

from phi.torch.flow import *
TORCH.set_default_device('GPU')

Suppose for example, later in my code, I'd like to track the learning history of some grids, so I'd like to move them to the CPU in order to save memory on my GPU. In native PyTorch, I'd call something like tensor.cpu() or tensor.to("cpu"). I'm wondering if there is an API in PhiFlow that does the same thing? I couldn't find it in the documentation, but I might be missing something here.

Looking at previous issues (#38), I came up with a workaround that looks something like the following, but it seems very cumbersome.

from phi.torch.flow import *
TORCH.set_default_device('GPU')

# Step 1 - Create a smoke grid on the GPU
test_smoke = CenteredGrid(0, extrapolation.BOUNDARY, x=resolution, y=resolution, bounds=Box(x=resolution, y=resolution))
print(test_smoke.values.native(order=test_smoke.shape).device)
# cuda:0

# Step 2 - Intermediate PyTorch tensor to move to the CPU
test_smoke_pt_tensor = test_smoke.values.native(order=test_smoke.shape).to("cpu")
print(test_smoke_pt_tensor.device)
# cpu

# Step 3 - Reconstruct the grid from the CPU tensor
test_smoke_cpu = test_smoke.with_values(tensor(test_smoke_pt_tensor, test_smoke.shape))
print(test_smoke_cpu.values.native(order=test_smoke.shape).device)
# cpu

Thanks for your help!

[holl-]

Device handling isn't developed very far in phiflow yet. Your method works. For a backend-indepenent way, you could write

data = data._op1(partial(data.default_backend.allocate_on_device, device=data.default_backend.list_devices('CPU')[0]))

where data is a tensor, e.g. test_smoke.values. I'll probably add a public function to do that for the next version.

[holl-]

I just pushed math.to_device() to 2.4-develop. Now you can write

test_smoke_cpu = math.to_device(test_smoke, 'CPU')

without having to access the values explicitly. Note that you will have to set convert=False when moving grids to the GPU, else it will try to convert the other tensors (grid size, extrapolation) as well.

[Dxyk]

This helps a lot. Thanks again!

[Dxyk]

Hi, this is a bit dated but here's a follow-up to this issue.

When converting to a PyTorch Tensor from a PhiFlow Tensor that was moved from the GPU to CPU using math.to_device(), it doesn't seem to function correctly.

TORCH.set_default_device("GPU")

smoke = CenteredGrid(0, extrapolation=extrapolation.BOUNDARY, bounds=Box(x=32, y=32), resolution=spatial(x=32, y=32))
print(smoke.values.device)
# GPU

smoke_cpu = math.to_device(smoke, "CPU", convert=False)
print(smoke_cpu.values.device)
# CPU

smoke_torch = smoke_cpu.values.native("x, y")
print(type(smoke_torch))
# <class 'torch.Tensor'>
print(smoke_torch.device)
# cuda:0 <-- Expected CPU

# Need to manually convert
smoke_torch_cpu = smoke_torch.to("cpu")
print(smoke_torch_cpu.device)
# cpu

[holl-]

@Dxyk Hi, right, there was a bug that led to allocation on the default device within native() for PyTorch. I now fixed this in phiml/develop. You'll need to install Φ-Flow 3.0 to get the update:

pip install --upgrade git+https://github.com/tum-pbs/PhiFlow@3.0