Tensor - N-dimensional arrays
Autograd - Automatic differentiation engine to compute gradients
Neural Network API - A modular framework for building Neural Networks with pre-defined layers, activation functions, and loss functions (nn.Module)
DataLoaders - Tools for efficient data handling. Allows for optimized training loops.

Syntax:

Imports:

torch: base
torch.nn: Neural Network model

import torch 
from torch import nn

Device / Cuda support:

device = "cuda" if torch.cuda.is_available() else "cpu"

Creating a model:

class LinearRegressionModel(nn.Module): # Pass nn.Module always
	def __init__(self):
		super().__init__() # Call initializer of parent class (nn.Module) 
		
		# Creates random weights
		self.weights = nn.Parameter(torch.randn(1, dtype = torch.float))
		
		# Creates random Bias
		self.bias = nn.Parameter(torch.randn(1, dtype = torch.float))
 
	# Define forward pass
	def forward(self, x: torch.Tensor) -> torch.Tensor:
		return self.weights * x + self.bias

Training Loop / Model Options

Loss Functions

nn.L1Loss() - MAE
nn.MSELoss() - MSE
nn.SmoothL1Loss() - Huber Loss
nn.CrossEntropyLoss() - Cross-Entropy Loss
nn.BCELoss() - Binary Cross-Entropy Loss
nn.HingeEmbeddingLoss() - Hinge Loss

loss_functionn = nn.L1Loss()

Optimizers

torch.optim.Adam() - Adam
torch.optim.AdamW() - AdamW
torch.optim.SGD() - SGD
torch.optim.RMSprop() - RMSprop
torch.optim.LBFGS() - L-BFGS

optimizer = torch.optim.SGD(params=model_0.parameters(), lr=.01)

Activation Functions

nn.ReLU() - ReLU
nn.LeakyReLU() - Leaky ReLU
nn.GELU() - GELU
nn.Sigmoid() - Sigmoid

relu = nn.ReLU()
layer = nn.Sequential( 
	nn.Linear(64, 32), 
	nn.ReLU() 
)

Layers

Linear Layers

nn.Linear() - Standard Fully Connected Layer
nn.Identity() - Identity (Pass-through)

nn.Linear(in_features=64, out_features=32)
nn.Identity()

Convolutional Layers

nn.Conv1d() - 1D Convolution
nn.Conv2d() - 2D Convolution
nn.Conv3d() - 3D Convolution

nn.Conv1d(in_channels=1, out_channels=16, kernel_size=3)
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1) 
nn.Conv3d(in_channels=3, out_channels=16, kernel_size=3)

Pooling Layers

nn.MaxPool2d() - Max Pooling
nn.AvgPool2d() - Average Pooling
nn.AdaptiveAvgPool2d() - Adaptive Pooling

nn.MaxPool2d(kernel_size=2, stride=2)
nn.AvgPool2d(kernel_size=2, stride=2)
nn.AdaptiveAvgPool2d(output_size=(1, 1))

Recurrent Layers

nn.RNN() - RNN
nn.LSTM() - LSTM
nn.GRU() - GRU

nn.RNN(input_size=10, hidden_size=20, batch_first=True) 
nn.LSTM(input_size=10, hidden_size=20, num_layers=2, batch_first=True) 
nn.GRU(input_size=10, hidden_size=20, batch_first=True)

Transformer Layers

nn.MultiheadAttention() - Multi-Head Attention
nn.TransformerEncoderLayer() - Transformer Encoder Layer

nn.MultiheadAttention(embed_dim=256, num_heads=8)
nn.TransformerEncoderLayer(d_model=256, nhead=8)

Regularization Layers

nn.Dropout() - Dropout
nn.Dropout2d() - Dropout

nn.Dropout(p=0.5) 
nn.Dropout2d(p=0.5) # For images (drops entire channels)

Normalization Layers

nn.BatchNorm1d() - Batch Normalization
nn.BatchNorm2d() - Batch Normalization for Images
nn.LayerNorm() - Layer Normalization
nn.InstanceNorm2d() - Instance Normalization

bn1d = nn.BatchNorm1d(num_features=32)
ln = nn.LayerNorm(normalized_shape=32)

Utility Layers

nn.Flatten() - Flatten
nn.Embedding() - Embedding

nn.Flatten()
nn.Embedding(num_embeddings=1000, embedding_dim=128)

Validation / Inference Loop

model.eval() #Sets model to evaluation mode (disables Dropout/BatchNorm).
with torch.no_grad(): #Disables gradient tracking (faster, saves memory)
	for data, target in dataloader: 
		data, target = data.to(device), target.to(device) 
		preds = model(data) # Calculate metrics (Accuracy/Loss) here

Learning Rate Schedulers

StepLR() - Step Decay
CosineAnnealingLR() - Cosine Annealing
ReduceLROnPlateau() - ReduceLROnPlateau

# 1. Standard (Step, Cosine, Exponential) 
for epoch in range(epochs): 
	train(...) 
	validate(...) 
	scheduler.step() # <-- Indented at EPOCH level
 
# 2. Reactive (ReduceLROnPlateau) 
for epoch in range(epochs): 
	train(...) 
	val_loss = validate(...) 
	scheduler.step(val_loss) # <-- Pass the metric here

Saving and Loading Models

# Saving torch.save
(model.state_dict(), "model_weights.pth") 
# Loading 
model = LinearRegressionModel() # Must initialize structure first 
model.load_state_dict(torch.load("model_weights.pth")) 
model.to(device)

Full Example Code

import torch
from torch import nn
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path
import save_model
 
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device: {device}")
weight = 0.7
bias = 0.3
 
start = 0
end = 1
step = 0.02
 
X = torch.arange(start, end, step).unsqueeze(dim=1)
y = weight * X + bias
 
train_split = int(0.8 * len(X))
X_train, y_train = X[:train_split], y[:train_split]
X_test, y_test = X[train_split:], y[train_split:]
 
def plot_predictions(train_data=X_train, 
                     train_labels=y_train, 
                     test_data=X_test, 
                     test_labels=y_test, 
                     predictions=None):
    plt.figure(figsize=(10, 7))
    plt.scatter(train_data, train_labels, c="b", s=4, label="Training data")
    plt.scatter(test_data, test_labels, c="g", s=4, label="Testing data")
    if predictions is not None:
        plt.scatter(test_data, predictions, c="r", s=4, label="Predictions")
    plt.legend(prop={"size": 14})
    plt.show()
 
class LinearRegressionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Parameter(torch.randn(1, 
                                               dtype = torch.float))
        self.bias = nn.Parameter(torch.randn(1,
                                            dtype = torch.float))
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.weights * x + self.bias
        
torch.manual_seed(42)        
 
model_0 = LinearRegressionModel()
 
# Loss Function
loss_fn = nn.L1Loss()
 
# Optimizer
optimizer = torch.optim.SGD(params=model_0.parameters(), lr=.01) 
 
epochs = 200
 
for epoch in range(epochs):
    # Set model to training mode
    model_0.train()
 
    # 1. Forward pass
    y_pred = model_0(X_train)
 
    # 2. Calculate loss
    loss = loss_fn(y_pred, y_train)
 
    # 3. Optimizer zero grad
    optimizer.zero_grad()
 
    # 4. Perform backpropogation on the loss with respect to parameters
    # of model
    loss.backward()
 
    # 5. Step the optimizer (perform gradient descent)
    optimizer.step()
 
    ### Testing
    model_0.eval()
    with torch.inference_mode():
        # 1. Forward pass
        test_pred = model_0(X_test)
 
        # 2. Calculate loss
        test_loss = loss_fn(test_pred,y_test)
 
    if epoch % 10 == 0:
        epoch_count.append(epoch)
        loss_values.append(loss)
        test_loss_values.append(test_loss)
        print(f"Epoch: {epoch} | Loss: {loss} | Test loss: {test_loss}")
 
        print(model_0.state_dict())
 
with torch.inference_mode():
    y_preds_new = model_0(X_test)
 
model_name = "linear_regression_model_0.pth"
save_model.save_model(model_0, model_name)⏎

PyTorch Geometric

Used for Graph Neural Networks (GNN).

from torch_geometric.nn import ...

Data Structure

from torch_geometric.data import Data
 
# 1. Node Features (x)
# Shape: [Num_Nodes, Num_Features]
# Example: 3 atoms, each with 5 chemical features (atomic mass, electronegativity, etc.)
x = torch.tensor([[1, ...], [1, ...], [6, ...]], dtype=torch.float)
 
# 2. Graph Connectivity (edge_index)
# Shape: [2, Num_Edges] -> COO Format (Coordinate List)
# Top row is Source Node, Bottom row is Target Node.
edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long)
 
# 3. Target Label (y)
# Shape: [1] (Solubility value)
y = torch.tensor([1.24], dtype=torch.float)
 
data = Data(x=x, edge_index=edge_index, y=y)

Layers

GCNConv() - GCNConv
GraphSAGE() - GraphSAGE
GATConv() - GATConv
GINEConv() - GINEConv

from torch_geometric.nn import GCNConv 
conv1 = GCNConv(in_channels=16, out_channels=32)

Complete GNN

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, global_mean_pool
 
class MoleculeGNN(torch.nn.Module):
    def __init__(self, num_node_features):
        super().__init__()
        # Layer 1: Message Passing
        self.conv1 = GCNConv(num_node_features, 64)
        # Layer 2: Message Passing
        self.conv2 = GCNConv(64, 64)
        # Layer 3: Message Passing
        self.conv3 = GCNConv(64, 64)
        
        # Output Head
        self.lin = torch.nn.Linear(64, 1) # Regression output (1 value)
 
    def forward(self, x, edge_index, batch):
        # 1. Obtain node embeddings
        x = self.conv1(x, edge_index)
        x = x.relu()
        x = self.conv2(x, edge_index)
        x = x.relu()
        x = self.conv3(x, edge_index)
 
        # 2. Readout layer (Pool atoms -> Molecule)
		# 'batch' vector tells the pool which atoms belong to which molecule
        x = global_mean_pool(x, batch) 
 
        # 3. Final Prediction
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.lin(x)
        return x

Rupan

PyTorch

Syntax:

Imports:

Device / Cuda support:

Creating a model:

Training Loop / Model Options

Loss Functions

Optimizers

Activation Functions

Layers

Linear Layers

Convolutional Layers

Pooling Layers

Recurrent Layers

Transformer Layers

Regularization Layers

Normalization Layers

Utility Layers

Validation / Inference Loop

Learning Rate Schedulers

Saving and Loading Models

Full Example Code

PyTorch Geometric

Data Structure

Layers

Complete GNN

Backlinks