DL 笔记1
使用GPU
raining on GPU
Just like how you transfer a Tensor onto the GPU, you transfer the neural net onto the GPU.
Let’s first define our device as the first visible cuda device if we have CUDA available:
device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)
Assuming that we are on a CUDA machine, this should print a CUDA device:
print(device)
Out:
cuda:0
The rest of this section assumes that device is a CUDA device.
Then these methods will recursively go over all modules and convert their parameters and buffers to CUDA tensors:
net.to(device)
Remember that you will have to send the inputs and targets at every step to the GPU too:
inputs, labels = data[0].to(device), data[1].to(device)
Why don’t I notice MASSIVE speedup compared to CPU? Because your network is really small