Pytorch Gpu
Pytorch using GPU to training
- Make input data and model to cuda device
-
for Model
net = Net() device = None if torch.cuda.is_available(): device = torch.device("cuda") net = net.to(device) -
for data
inputs = inputs.to(device) labels = labels.to(device) -
verify
input.device # cuda label.device # cuda
-
- Make output data to cpu device
- for data
output = net(images) # images>gpu, net>gpu, output>gpu output = output.cpu() # output>cpu images = images.cpu() # images>cpu labels = labels.cpu() # labels>cpu
Follow [2] to training by pytorch
