PyTorch’s to method returns a tensor with the requested device, data type, or both. Use it to copy CPU data to GPU memory, move results back to the CPU, or cast values before computation. The tensor keeps the same numerical data when only its device changes, while subsequent operations remain on the selected device. section 3.10.1 section 3.6.3
Common forms include:
- points_gpu = points.to(device='cuda') copies a tensor to the GPU; device='cuda:0' selects the first GPU, and device='cpu' copies it back. section 3.10.1
- tensor.to(torch.double) or tensor.to(dtype=torch.short) casts a tensor; to checks whether conversion is needed before performing it. section 3.6.3
- Tensor.to is out of place: the original tensor remains unchanged and the method returns a new tensor when a conversion is required. section 8.4.3
- nn.Module.to moves a model’s parameters and can cast their type, but modifies the module in place; create the optimizer afterward. section 8.4.3 section 13.3.1
Device-aware training
Select an available device, move the model before constructing its optimizer, and move every input batch and label before the forward pass: device = torch.device('cuda' if torch.cuda.is_available() else 'cpu'); model = Net().to(device=device); imgs = imgs.to(device=device); labels = labels.to(device=device). section 8.4.3 section 13.3.1
# In[37]:
import datetime
def training_loop(n_epochs, optimizer, model, loss_fn, train_loader):
for epoch in range(1, n_epochs + 1):
start_time = datetime.datetime.now()
loss_train = 0.0
for imgs, labels in train_loader:
imgs = imgs.to(device=device) #1
labels = labels.to(device=device)
outputs = model(imgs)
loss = loss_fn(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_train += loss.item()
end_time = datetime.datetime.now()
epoch_duration = (end_time - start_time).total_seconds()
if epoch == 1 or epoch % 10 == 0:
print('{} Epoch {}, Training loss {:.6f}, Time {:.2f}s'.format(
datetime.datetime.now(), epoch,
loss_train / len(train_loader), epoch_duration))
Shows the two Tensor.to calls that place training images and labels on the selected device inside a training loop. section 8.4.3
