Saving and loading models (checkpoints)

3.3.7 · AI-ML › Deep Learning Frameworks

Core Problem: Hume Actually Kya Save Karna Hai?

Jab hum ek neural network train karte hain, toh hum gradient descent ke zariye lakho (ya arbon) parameters update kar rahe hote hain. Training state mein yeh sab aata hai:

1. Model Parameters: Har layer ke liye weights $\mathbf{W}$ aur biases $\mathbf{b}$
2. Optimizer State: Momentum buffers (SGD with momentum), running averages (Adam ka $m_t$, $v_t$), learning rate schedules
3. Epoch/Step Counter: Hum training mein kahan hain
4. Random Seeds: Reproducibility ke liye
5. Model Architecture (optional): Network ko reconstruct karne ka blueprint

Yeh sab kyun? Sirf parameters model ki predictions define karte hain, lekin training smoothly continue karne ke liye optimizer state zaroori hai. Agar aap sirf weights save karte ho aur baad mein fresh optimizer ke saath resume karte ho, toh accumulated momentum kho jaata hai—jaise koi gaadi cruising speed ke bajaay achanak zero velocity se shuru ho.

Framework-Specific Approaches

PyTorch: State Dictionaries

Yeh design kyun? Architecture (code) aur parameters (data) ka separation. Aapki model class evolve ho sakti hai, lekin purane checkpoints loadable rehte hain agar layer names match karein.

import torch
import torch.nn as nn
 
# Define model
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
# After training for N epochs...
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': train_loss,
    'accuracy': val_acc
}
torch.save(checkpoint, 'model_epoch_10.pt')

Yeh step kyun? Hum training resume karne ke liye zaroori sab kuch ek single dict mein bundle karte hain, phir torch.save se serialize karte hain (jo hood ke neeche Python ka pickle use karta hai).

Wapas load karna:

checkpoint = torch.load('model_epoch_10.pt')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch'] + 1

load_state_dict kyun aur direct assignment kyun nahi? Yeh tensor shapes aur missing/unexpected keys validate karta hai, architecture mismatches pakadta hai.

TensorFlow/Keras: Do Paradigms

1. TensorFlow Checkpoints (.ckpt)

Sirf weights ko binary files + ek index ke roop mein save karta hai:

model = tf.keras.Sequential([...])
checkpoint_path = "training_1/cp-{epoch:04d}.ckpt"
checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_path,
    save_weights_only=True,
    save_freq='epoch'
)
model.fit(x_train, y_train, epochs=10, callbacks=[checkpoint_callback])
 
# Load
model.load_weights('training_1/cp-010.ckpt')

save_weights_only=True kyun? Faster saves, chhoti files. Iske liye aapko model architecture alag se define karni padti hai.

2. SavedModel Format (Full Model)

model.save('my_model')  # Creates a directory with architecture + weights
loaded_model = tf.keras.models.load_model('my_model')

Yeh format kyun? Architecture, weights, aur computational graph tak ko bundle karta hai—original code ki zaroorat nahi. TensorFlow Serving ke saath production deployment ke liye ideal hai.

Checkpoint Strategies: Kab aur Kya Save Karein

Strategy 1: Sirf Best Model Save Karo

# PyTorch example
best_acc = 0.0
for epoch in range(num_epochs):
    val_acc = evaluate(model, val_loader)
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'accuracy': val_acc
        }, 'best_model.pt')

Yeh kyun kaam karta hai? Zyaadatar experiments ko sirf peak performance model chahiye hota hai. Intermediate checkpoints wasted storage hote hain.

Strategy 2: Har N Epochs Par Save Karo + Last K Rakhho

# Keras callback
checkpoint_callback = ModelCheckpoint(
    filepath='model_{epoch:02d}_{val_loss:.2f}.h5',
    save_best_only=False,
    save_freq='epoch',
    period=5  # Every 5 epochs
)

Phir manually purane checkpoints hatao ya rolling window use karo (sirf last 3 rakhho).

Kyun? Long training runs mein fallback points helpful hote hain. Agar epoch 47 par validation loss spike kar jaaye, toh aap epoch 45 se reload kar sakte ho.

Strategy 3: Plateau Par ya Early Stopping Par Save Karo

from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
 
early_stop = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)
checkpoint = ModelCheckpoint('checkpoint.h5', monitor='val_loss', save_best_only=True)

patience=10 kyun? Validation metrics noisy hoti hain. Ek bura epoch training rokne ka trigger nahi hona chahiye. Patience model ko local minima se nikalne ka waqt deta hai.

Derivation: Optimizer State Kyun Mayne Rakhti Hai

SGD with momentum consider karo:

Jahan:

• ${\mathbf{v}}_t$ velocity hai (gradients ka moving average)
• $\beta$ momentum coefficient hai (typically 0.9)
• $\alpha$ learning rate hai

Agar aap sirf ${\theta }_t$ save karo aur restart karo toh kya hoga?

Resume par, ${\mathbf{v}}_t$ $0$ par reset ho jaata hai. Pehla step vanilla SGD ban jaata hai:

${\theta }_{t+1}={\theta }_t-\alpha {\nabla }_{\theta }\mathcal{L}$

Yeh bura kyun hai? Momentum kai steps mein gradient information accumulate karta hai. Isse reset karne ka matlab hai:

1. Directional inertia ka loss: Model "bhool jaata hai" ki kaunsi direction promising thi
2. Slower convergence: Momentum rebuild karne mein ~$1/(1-\beta )$ steps lagte hain ($\beta =0.9$ ke liye 11 steps)
3. Unstable training: Agar aap kisi sharp valley mein ho, toh pehle kuch steps diverge ho sakte hain

Fix: optimizer.state_dict() save karo jo har parameter ke liye ${\mathbf{v}}_t$ include karta hai.

Example 1: Crash Ke Baad Training Resume Karna

import torch
import torch.nn as nn
 
# Training loop with checkpointing
def train_with_checkpoints(model, optimizer, train_loader, num_epochs, checkpoint_dir):
    start_epoch = 0
    # Try to resume from latest checkpoint
    checkpoint_files = sorted(Path(checkpoint_dir).glob('checkpoint_*.pt'))
    if checkpoint_files:
        latest = checkpoint_files[-1]
        checkpoint = torch.load(latest)
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        start_epoch = checkpoint['epoch'] + 1
        print(f"Resumed from epoch {start_epoch}")
    for epoch in range(start_epoch, num_epochs):
        for batch in train_loader:
            loss = train_step(model, batch, optimizer)
        # Save checkpoint every epoch
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss.item()
        }, f'{checkpoint_dir}/checkpoint_{epoch:04d}.pt')

start_epoch adjustment kyun? Kaam repeat karne se bachata hai. Agar epoch 7 par crash hua, toh hum epoch 8 se shuru karte hain.

Sort kyun karo aur [-1] kyun lo? Files out of order list ho sakti hain. Naam se sort karna (zero-padded numbers ke saath) ensure karta hai ki hum truly latest checkpoint load karein.

Example 2: Inference Ke Liye Export Karna (Optimizer Ke Bina)

# Training complete, save inference-only checkpoint
torch.save({
    'model_state_dict': model.state_dict(),
    'accuracy': final_test_acc,
    'class_names': ['cat', 'dog', 'bird']
}, 'production_model.pt')
 
# Load in production
checkpoint = torch.load('production_model.pt', map_location='cpu')
model = MyModel()  # Architecture must match
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()  # Disable dropout/batchnorm training mode
 
# Inference
with torch.no_grad():
    predictions = model(new_data)

map_location='cpu' kyun? Model GPU (CUDA) par train hua tha. Bina is flag ke directly CPU server par load karna fail ho jaata.

model.eval() kyun? Dropout aur BatchNorm training (stochastic) ke dauraan alag behave karte hain compared to inference (deterministic) ke dauraan. Ise bhoolna ek common deployment bug hai.

Advanced: Distributed Training Checkpoints

Jab multiple GPUs par training hoti hai, har device model ka ek shard hold karta hai. Checkpointing mein coordination zaroori hai:

# PyTorch DistributedDataParallel
if torch.distributed.get_rank() == 0:  # Only master process saves
    torch.save({
        'model_state_dict': model.module.state_dict(),  # .module unwraps DP
        'optimizer_state_dict': optimizer.state_dict()
    }, 'checkpoint.pt')

Sirf rank 0 kyun? Saare GPUs ke parameters identical hote hain (AllReduce ke zariye synchronized). Redundantly save karna I/O waste karta hai.

.module kyun? DP aapke model ko ek DistributedDataParallel object mein wrap karta hai. Actual model model.module hai.

Connections

• Overfitting and Regularization: Checkpoints early stopping enable karte hain validation loss diverge hone se pehle model save karke.
• Learning Rate Schedules: LR schedulers epoch counters par depend karte hain; checkpoints ko yeh state preserve karni chahiye.
• Distributed Training: Multi-GPU setups mein data corruption se bachne ke liye synchronized checkpointing zaroori hai.
• Model Deployment: Production models checkpoints se load hote hain; architecture mismatches runtime failures karti hain.
• Transfer Learning: Pretrained models (BERT, ResNet) sirf kisi aur ki training run ke checkpoints hote hain.

#flashcards/ai-ml

Deep learning mein checkpoint kya hota hai? :: Ek model ke parameters (weights, biases), optimizer state, aur training metadata (epoch, loss) ka disk par saved snapshot, jo training resumption ya deployment allow karta hai.

Concept Map