How to specify CPU and GPU for training in pytorch, and switch between cpu and gpu
I always don't understand things about cpu and gpu in pytorch. Recently, it has become clearer. Write it down and record it for future reference.
Method 1: .to(device)
Recommended use~
1. When you do not know whether the computer gpu is available or not:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu' )
a.to(device)
The first line of code means to determine whether the computer GPU is available. If it is available, the device uses cuda() to call the GPU. If it is not available, it uses cpu() to call the CPU.
The second line of code means to put the variable on the corresponding device (of course, if you are using the CPU, you don't need this step, because the variable is stored on the CPU by default. If you call the GPU, you must first put the variable on the GPU. Run, and then call back to the CPU after running)
2. When specifying gpu :
device = torch.device('cuda')
2. When specifying cpu :
device = torch.device('cpu')
Method 2: .cuda() & .cpu()
The above method only needs to modify the device if you want to switch between CPU and GPU, and this method is equivalent to manually transferring variables to CPU or GPU, which is troublesome and not recommended . Here is just to demonstrate the usage:
1. Specify When gpu :
a.cuda()
2. When specifying cpu :
a.cpu()
Method 3: os.environ["CUDA_VISIBLE_DEVICES"]
In fact, it is not a method, CUDA_VISIBLE_DEVICESit means that the current graphics card can be detected by the python environment program, so when calling the GPU, adding this line of code means specifying which graphics card to use for the GPU
1. When using gpu :
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
Adding this line of code to the front of the program means that when calling the GPU, the 0th graphics card is used (if you want to use the 1st block, change the number 0 to 1)
2. When using the cpu :
os.environ["CUDA_VISIBLE_DEVICES"] = ' '
Just change the number to a space to indicate that the CPU is used (because the space indicates that the graphics card cannot be detected, so it is the CPU)