Part 8: Tensor Operations with GPU Acceleration.

One of the significant advantages of using PyTorch is its seamless integration with GPUs, which can significantly accelerate tensor operations and model training. PyTorch provides simple APIs to move tensors and models to a GPU.

Checking for GPU Availability

Before performing tensor operations on a GPU, it is essential to check if a GPU is available.

Check for GPU:

import torch

if torch.cuda.is_available():
    device = torch.device("cuda")
    print("CUDA is available. Using GPU.")
else:
    device = torch.device("cpu")
    print("CUDA is not available. Using CPU.")

Moving Tensors to GPU

You can move tensors to a GPU using the to method or the cuda method.

Move a Tensor to GPU:

tensor = torch.tensor([1, 2, 3, 4, 5])
tensor_gpu = tensor.to(device)
print(tensor_gpu)

Creating a Tensor Directly on GPU:

tensor_gpu = torch.tensor([1, 2, 3, 4, 5], device=device)
print(tensor_gpu)

Performing Operations on GPU

Once tensors are on the GPU, you can perform operations on them just like you would on the CPU.

Basic Operations on GPU:

tensor_a = torch.tensor([1, 2, 3], device=device)
tensor_b = torch.tensor([4, 5, 6], device=device)

tensor_sum = tensor_a + tensor_b
print(tensor_sum)

tensor_prod = tensor_a * tensor_b
print(tensor_prod)

Example: Neural Network Training on GPU

Moving a neural network and its tensors to the GPU can significantly speed up training.

Define a Simple Neural Network:

import torch.nn as nn

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(10, 50)
        self.fc2 = nn.Linear(50, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleNN().to(device)

Training Loop on GPU:

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()

inputs = torch.randn(64, 10, device=device)
targets = torch.randn(64, 1, device=device)

for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()

print("Training completed.")

In this example, both the inputs and the model are moved to the GPU, and the training loop performs all operations on the GPU, resulting in faster computation.