PyTorch: Infer dtype from device capability, not input data

Question

I am trying to write a device-agnostic library for PyTorch, and I have stumbled across the problem of PyTorch using dtypes that are not compatible with my compute device:

import scipy.signal
import torch

raw_window = scipy.signal.windows.cosine(128)
print(raw_window.dtype)  # float64

device = torch.device("cpu")
window = torch.as_tensor(raw_window, device=device)
print(window.device)  # cpu
print(window.dtype)   # torch.float64

device = torch.device("cuda")
window = torch.as_tensor(raw_window, device=device)
print(window.device)  # cuda:0
print(window.dtype)   # torch.float64

As you can see in the last line, torch assigns the dtype torch.float64, even though my CUDA device is not able to handle double precision float values.

Is there a way to make PyTorch use the most suitable device dtype instead of the input data dtype? Or am I misunderstanding that whole concept of setting devices and dtypes completely?

score 0 · Answer 1 · answered Apr 15 '20 at 11:40

No, as you noticed PyTorch infers dtype from input data only.

In your case, as numpy has it's default set to np.float64 (regardless of system and architecture) PyTorch will infer it's analogous torch.float64, so it's more of a problem with starting from numpy (and you can't set different default dtype).

In pytorch you usually go for torch.float32 (and it is the default), eventually torch.float16 for mixed precision (leaving out quantization). Usually such precision is sufficient (not sure about your exact use-case though).

Your best bet is hence to cast numpy arrays to np.float32 like this:

raw_window = raw_window.astype(np.float32)

(or cast it via pytorch or even better create float32 from the beginning), which AFAIK should be available on all/most devices be it CPU or GPU.

For CPU/GPU agnostic approach you should go for explicit device creation as shown here:

if not args.disable_cuda and torch.cuda.is_available():
    args.device = torch.device('cuda')
else:
    args.device = torch.device('cpu')

(extend if you need to explicitly specify cuda or a-like) and use it during tensor creation.

Your last remark is a little bit redundant as that is exactly what I am doing already. — Nils Werner, Apr 16 '20 at 07:39
@NilsWerner So i thought, but better redundant and too explicit than too implicit — Szymon Maszke, Apr 16 '20 at 09:28

PyTorch: Infer dtype from device capability, not input data

1 Answers1