How to use the device function (pytorch)

TOP MENU

Python

library
・pip

・MeCab

numpy
・digitize

・mgrid

・pad

・polyfit

・prod

・shape

matplotlib
・figure

・pcolormesh

・scatter

pytorch
・BCELoss, MSELoss

・device

・Embedding

・TensorDataset, Dataloader

・RNN, LSTM

scikit-learn
・SVC

scipy
・interpolate

tkinter
・postscript

・image display

・frame, grid

other
・linear interpolation

OpenAI gym
・CartPole-v0

Release date:2022/10/20　　　　　　　　　

・In Japanese
＜premise knowledge＞
・Python

■Description
Specify the GPU (cuda) to use with pytorch. cuda is a development environment for GPU provided by NVIDIA.

■Example
First, check if your PC has a GPU that can be used with pytorch.

import torch
print(torch.cuda.is_available())

⇒ False # False if there is no cuda environment

Next, store either "cuda" or "cpu" in a variable as a device name that can be used, as follows.

dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") #(1)
print(dev)

⇒ device(type='cpu')

By using ".to", you can perform operations using cuda or cpu as follows.

b = torch.zeros(4)
c = b.to(dev)
print(c)

⇒ tensor([0., 0., 0., 0.])

If you get an error message like the one below, the GPU is set to be used even though it cannot be used. Let's check if the formula in #(1) above is set correctly.

　"Torch not compiled with CUDA enabled"

List of related articles