DEV Community

Super Kai (Kazuya Ito)
Super Kai (Kazuya Ito)

Posted on • Edited on

Set and get `device` in PyTorch

#python #pytorch #device #argument

Buy Me a Coffee☕

*Memos:

You can set and get device as shown below:

*Memos:

  • I selected some popular device argument functions such as tensor(), arange(), rand(), rand_like(), zeros() and zeros_like().
  • Basically, device(Optional-Default:None-Type:int, str or device()).
  • Basically, if device is None, it's inferred from other tensor or get_default_device() is used. *My post explains get_default_device() and set_default_device().
  • cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia or privateuseone can be set to device.
  • Setting 0 to device uses cuda(GPU). *The number must be zero or positive.
  • Basically, device= must be needed.
  • My post explains device().
  • str() can get a device value.

tensor(). *My post explains tensor():

import torch my_tensor = torch.tensor([0, 1, 2]) my_tensor = torch.tensor([0, 1, 2], device='cpu') my_tensor = torch.tensor([0, 1, 2], device=torch.device(device='cpu')) my_tensor = torch.tensor([0, 1, 2], device=torch.device(type='cpu')) my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([0, 1, 2]), device(type='cpu'), 'cpu')  my_tensor = torch.tensor([0, 1, 2], device='cuda:0') my_tensor = torch.tensor([0, 1, 2], device='cuda') my_tensor = torch.tensor([0, 1, 2], device=0) my_tensor = torch.tensor([0, 1, 2], device=torch.device(device='cuda:0')) my_tensor = torch.tensor([0, 1, 2], device=torch.device(device='cuda')) my_tensor = torch.tensor([0, 1, 2], device=torch.device(device=0)) my_tensor = torch.tensor([0, 1, 2], device=torch.device(type='cuda', index=0)) my_tensor = torch.tensor([0, 1, 2], device=torch.device(type='cuda')) my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([0, 1, 2], device='cuda:0'), device(type='cuda', index=0), 'cuda:0')
Enter fullscreen mode Exit fullscreen mode

tensor() with is_available(). *My post explains is_available():

import torch my_device = "cuda:0" if torch.cuda.is_available() else "cpu" my_tensor = torch.tensor([0, 1, 2], device=my_device) my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([0, 1, 2], device='cuda:0'), device(type='cuda', index=0), 'cuda:0')
Enter fullscreen mode Exit fullscreen mode

arange(). *My post explains arange():

import torch my_tensor = torch.arange(start=5, end=15, step=3, device='cpu') my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([5, 8, 11, 14]), device(type='cpu'), 'cpu')  my_tensor = torch.arange(start=5, end=15, step=3, device='cuda:0') my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([5, 8, 11, 14], device='cuda:0'), # device(type='cuda', index=0), # 'cuda:0')
Enter fullscreen mode Exit fullscreen mode

rand(). *My post explains rand():

import torch my_tensor = torch.rand(size=(3,), device='cpu') my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([0.2782, 0.3780, 0.6509]), device(type='cpu'), 'cpu')  my_tensor = torch.rand(size=(3,), device='cuda:0') my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([0.1052, 0.9281, 0.0151], device='cuda:0'), # device(type='cuda', index=0), # 'cuda:0')
Enter fullscreen mode Exit fullscreen mode

rand_like(). *My post explains rand_like():

import torch my_tensor = torch.rand_like(input=torch.tensor([7., 4., 5.]), device='cpu') my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([0.9130, 0.7072, 0.1935]), device(type='cpu'), 'cpu')  my_tensor = torch.rand_like(input=torch.tensor([7., 4., 5.]), device='cuda:0') my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([0.3655, 0.6319, 0.3045], device='cuda:0'), # device(type='cuda', index=0), # 'cuda:0')
Enter fullscreen mode Exit fullscreen mode

zeros(). *My post explains zeros():

import torch my_tensor = torch.zeros(size=(3,), device='cpu') my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([0., 0., 0.]), device(type='cpu'), 'cpu')  my_tensor = torch.zeros(size=(3,), device='cuda:0') my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([0., 0., 0.], device='cuda:0'), # device(type='cuda', index=0), # 'cuda:0')
Enter fullscreen mode Exit fullscreen mode

zeros_like(). *My post explains zeros_like():

import torch my_tensor = torch.zeros_like(input=torch.tensor([7., 4., 5.]), device='cpu') my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([0., 0., 0.]), device(type='cpu'), 'cpu')  my_tensor = torch.zeros_like(input=torch.tensor([7., 4., 5.]), device='cuda:0') my_tensor, my_tensor.device, str(my_tensor.device) # (tensor([0., 0., 0.], device='cuda:0'), # device(type='cuda', index=0), # 'cuda:0')
Enter fullscreen mode Exit fullscreen mode

Top comments (0)