Practical: PyTorch

# Installation
!pip3 install torch torchvision torchaudio torch_geometric

import torch
import numpy as np

# --- Creating Tensors ---

# From a list
data = [[1, 2], [3, 4]]
x_data = torch.tensor(data)
print(f"Tensor from list:\n{x_data}\n")

# From a NumPy array (shares memory by default!)
np_array = np.array(data)
x_np = torch.from_numpy(np_array)
print(f"Tensor from NumPy array:\n{x_np}\n")

# Creating tensors with specific shapes and values
shape = (2, 3,)
rand_tensor = torch.rand(shape) # Random values between 0 and 1
ones_tensor = torch.ones(shape)   # All ones
zeros_tensor = torch.zeros(shape)  # All zeros
print(f"Random Tensor:\n {rand_tensor} \n")
print(f"Ones Tensor:\n {ones_tensor} \n")
print(f"Zeros Tensor:\n {zeros_tensor}\n")

# --- Tensor Attributes ---
tensor = torch.rand(3, 4)
print(f"Shape of tensor: {tensor.shape}")
print(f"Datatype of tensor: {tensor.dtype}")
print(f"Device tensor is stored on: {tensor.device}") # Default is CPU

# --- Moving Tensors to GPU (if available) ---
if torch.cuda.is_available():
    print("CUDA (GPU) is available!")
    device = torch.device("cuda")
    # Move tensor to GPU
    tensor_gpu = tensor.to(device)
    print(f"Device tensor_gpu is stored on: {tensor_gpu.device}")
    # Operations on tensor_gpu will run on the GPU

    # Move back to CPU (e.g., for use with NumPy)
    tensor_cpu = tensor_gpu.to("cpu")
    # NOTE: NumPy cannot handle GPU tensors directly.

elif torch.backends.mps.is_available():
    print("Apple Silicon GPU (MPS) is available!")
    device = torch.device("mps")
    # Move tensor to MPS
    tensor_mps = tensor.to(device)
    print(f"Device tensor_mps is stored on: {tensor_mps.device}")
    # Operations on tensor_mps will run on the Apple Silicon GPU

    # Move back to CPU (e.g., for use with NumPy)
    tensor_cpu = tensor_mps.to("cpu")
    # NOTE: NumPy cannot handle MPS tensors directly.
else:
    print("CUDA or Apple Silicon (GPU) not available, using CPU.")
    device = torch.device("cpu") # Use CPU if GPU not available

# --- Basic Operations ---
# Similar syntax to NumPy
tensor_a = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
tensor_b = torch.tensor([[5, 6], [7, 8]], dtype=torch.float32)

# Element-wise addition
sum_tensor = tensor_a + tensor_b
# or torch.add(tensor_a, tensor_b)
print(f"Addition:\n{sum_tensor}\n")

# Element-wise multiplication
mul_tensor = tensor_a * tensor_b
# or torch.mul(tensor_a, tensor_b)
print(f"Element-wise Multiplication:\n{mul_tensor}\n")

# Matrix multiplication
matmul_tensor = tensor_a @ tensor_b
# or torch.matmul(tensor_a, tensor_b)
print(f"Matrix Multiplication:\n{matmul_tensor}\n")

# --- NumPy Bridge ---
# Tensor to NumPy array
numpy_array_again = sum_tensor.numpy() # Only works for CPU tensors
print(f"Tensor converted back to NumPy:\n{numpy_array_again}\n")

# Create tensors that require gradient tracking
x = torch.ones(2, 2, requires_grad=True)
print(f"x:\n{x}\n")

# Perform an operation
y = x + 2
print(f"y = x + 2:\n{y}\n")
# y was created as a result of an operation involving x, so it has a grad_fn
print(f"y.grad_fn: {y.grad_fn}\n")

z = y * y * 3
out = z.mean() # Calculate a scalar value
print(f"z = y*y*3:\n{z}\n")
print(f"out = z.mean(): {out}\n")

# --- Compute gradients ---
# out is a scalar, so we can call backward() directly
out.backward()

# Gradients are accumulated in the .grad attribute of the tensors
# d(out)/dx is computed
print(f"Gradient of out w.r.t. x (d(out)/dx):\n{x.grad}\n")

# Let's verify manually for one element:
# out = (1/4) * sum(z) = (1/4) * sum(3 * (x+2)^2)
# d(out)/dx_ij = (1/4) * d/dx_ij [ 3 * (x_ij+2)^2 ]
#             = (1/4) * 3 * 2 * (x_ij+2) * 1
#             = (3/2) * (x_ij+2)
# Since x starts as ones(2,2), x_ij = 1.
# d(out)/dx_ij = (3/2) * (1+2) = (3/2) * 3 = 4.5
# The output tensor x.grad should be [[4.5, 4.5], [4.5, 4.5]]

# Turn off gradient tracking when not needed (e.g., during evaluation)
with torch.no_grad():
    y_no_grad = x + 2
    print(f"y_no_grad.requires_grad: {y_no_grad.requires_grad}")

import torch.nn as nn

# Example: A simple Linear Regression Model (y = Wx + b)
class LinearRegression(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LinearRegression, self).__init__()
        # Define the layer(s)
        self.linear = nn.Linear(input_dim, output_dim) # One linear layer

    def forward(self, x):
        # Define the forward pass
        out = self.linear(x)
        return out

# Example: A Multi-Layer Perceptron (like used in the demos)
class SimpleMLP(nn.Module):
    def __init__(self, input_size, output_size=1, hidden_layers=[64, 32]):
        super(SimpleMLP, self).__init__()
        layers = []
        current_size = input_size
        # Dynamically create hidden layers based on the list
        for hidden_size in hidden_layers:
            layers.append(nn.Linear(current_size, hidden_size))
            layers.append(nn.ReLU()) # Add ReLU activation after each hidden layer
            current_size = hidden_size
        # Add the final output layer
        layers.append(nn.Linear(current_size, output_size))
        # Use nn.Sequential to easily chain the layers
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        # Pass input through the sequential network
        return self.network(x)

# Instantiate the MLP model
input_features = 10 # Example input size
output_value = 1    # Example output size (for regression)
model_mlp = SimpleMLP(input_features, output_value)
print("MLP Model Structure:")
print(model_mlp)

# You can easily inspect model parameters
print("\nModel Parameters:")
for name, param in model_mlp.named_parameters():
    if param.requires_grad:
        print(name, param.shape)

from torch import nn
import torch

# Example Loss Functions
criterion_mse = nn.MSELoss() # Mean Squared Error: For regression
criterion_bce = nn.BCELoss() # Binary Cross Entropy: For binary classification (with Sigmoid output)
criterion_ce = nn.CrossEntropyLoss() # Cross Entropy Loss: For multi-class classification (expects raw logits)

# Example usage (assuming model_output and target are tensors)
model_output_reg = torch.randn(10, 1, requires_grad=True) # Example regression output
target_reg = torch.randn(10, 1)
loss_mse = criterion_mse(model_output_reg, target_reg)
print(f"Example MSE Loss: {loss_mse.item()}") # .item() gets scalar value

model_output_cls = torch.sigmoid(torch.randn(10, 1, requires_grad=True)) # Example classification output (post-sigmoid)
target_cls = torch.randint(0, 2, (10, 1)).float() # Example binary targets (0 or 1)
loss_bce = criterion_bce(model_output_cls, target_cls)
print(f"Example BCE Loss: {loss_bce.item()}")

import torch.optim as optim

# Use the MLP model we defined earlier
model_to_train = SimpleMLP(10, 1)

# --- Choose an Optimizer ---
# Adam is a popular and often effective choice
optimizer_adam = optim.Adam(model_to_train.parameters(), lr=0.001) # lr is the learning rate

# Stochastic Gradient Descent (SGD) is another common one
# optimizer_sgd = optim.SGD(model_to_train.parameters(), lr=0.01, momentum=0.9)

# --- Optimizer Steps (within a training loop) ---
# 1. Zero the gradients before calculating loss for a new batch
#    optimizer_adam.zero_grad()
#
# 2. Calculate the loss
#    loss = criterion(outputs, targets)
#
# 3. Compute gradients w.r.t. parameters
#    loss.backward()
#
# 4. Update the model parameters based on gradients
#    optimizer_adam.step()

from torch.utils.data import TensorDataset, DataLoader

# Example data (already tensors)
features = torch.randn(100, 10) # 100 samples, 10 features each
labels = torch.randn(100, 1)   # 100 corresponding labels

# Create a TensorDataset
dataset = TensorDataset(features, labels)

# Access a single sample
first_sample_features, first_sample_label = dataset[0]
print(f"First sample features shape: {first_sample_features.shape}")
print(f"First sample label shape: {first_sample_label.shape}")

# Create a DataLoader
batch_size = 16
# shuffle=True is important for training to mix data between epochs
# num_workers > 0 can speed up loading using parallel processes (use cautiously in notebooks)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)

# Iterate over the DataLoader to get batches
print("\nIterating through DataLoader:")
for batch_idx, (batch_features, batch_labels) in enumerate(dataloader):
    print(f"Batch {batch_idx+1}:")
    print(f"  Features shape: {batch_features.shape}") # Should be [batch_size, 10]
    print(f"  Labels shape: {batch_labels.shape}")     # Should be [batch_size, 1]
    # --- Inside the training loop, you'd process this batch ---
    # model(batch_features) -> calculate loss -> loss.backward() -> optimizer.step()
    if batch_idx >= 2: # Stop after showing a few batches
      break

# --- Setup (Assume model, criterion, optimizer, train_loader are defined) ---
# model = YourModel(...)
# criterion = YourLossFunction()
# optimizer = YourOptimizer(model.parameters(), lr=...)
# train_loader = DataLoader(your_train_dataset, ...)
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# model.to(device) # Move model to the appropriate device

# Let's create dummy versions for illustration:
model = LinearRegression(5, 1).to(device) # Simple model: 5 features -> 1 output
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
dummy_features = torch.randn(100, 5)
dummy_labels = torch.randn(100, 1) * 3 + 2 # y = 3x + 2 + noise approx
dummy_dataset = TensorDataset(dummy_features, dummy_labels)
train_loader = DataLoader(dummy_dataset, batch_size=10, shuffle=True)

# --- Training Loop ---
num_epochs = 10
print("\nStarting Dummy Training Loop:")
model.train() # Set the model to training mode (important for layers like dropout, batchnorm)

for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (inputs, targets) in enumerate(train_loader):
        # Move data to the same device as the model
        inputs = inputs.to(device)
        targets = targets.to(device)

        # 1. Zero the parameter gradients
        optimizer.zero_grad()

        # 2. Forward pass: compute predicted outputs
        outputs = model(inputs)

        # 3. Calculate the loss
        loss = criterion(outputs, targets)

        # 4. Backward pass: compute gradient of the loss w.r.t. parameters
        loss.backward()

        # 5. Perform a single optimization step (parameter update)
        optimizer.step()

        # Accumulate loss for reporting
        running_loss += loss.item() * inputs.size(0) # loss.item() gets scalar loss

    epoch_loss = running_loss / len(train_loader.dataset)
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}')

print('Finished Training')

# After training, you typically evaluate on a separate test set
# model.eval() # Set model to evaluation mode
# with torch.no_grad(): # Disable gradient calculation for evaluation
#    # ... loop through test_loader and calculate metrics ...