Training Time Optimization

1from torch.cuda.amp import autocast, GradScaler
2
3scaler = GradScaler()
4
5for batch in dataloader:
6    optimizer.zero_grad()
7
8    # Forward pass in FP16
9    with autocast():
10        rul_pred, health_pred = model(batch)
11        loss = amnl_loss(rul_pred, health_pred, targets)
12
13    # Backward pass with scaled gradients
14    scaler.scale(loss).backward()
15    scaler.step(optimizer)
16    scaler.update()
17
18# Speedup: ~1.4x (90 min → ~65 min)

1from torch.utils.data import DataLoader
2
3# Optimized DataLoader configuration
4train_loader = DataLoader(
5    train_dataset,
6    batch_size=256,
7    shuffle=True,
8    num_workers=8,       # Parallel data loading
9    pin_memory=True,     # Faster GPU transfer
10    prefetch_factor=2,   # Prefetch batches
11    persistent_workers=True  # Keep workers alive
12)
13
14# Speedup: ~1.15x (8 min data loading → ~4 min)

1accumulation_steps = 4
2effective_batch_size = 256 * 4  # = 1024
3
4for i, batch in enumerate(dataloader):
5    loss = compute_loss(batch) / accumulation_steps
6    loss.backward()
7
8    if (i + 1) % accumulation_steps == 0:
9        optimizer.step()
10        optimizer.zero_grad()
11
12# Benefit: Larger effective batch → faster convergence
13# ~10% fewer epochs needed for same performance

1class EarlyStopping:
2    def __init__(self, patience=15, min_delta=0.001):
3        self.patience = patience
4        self.min_delta = min_delta
5        self.best_loss = float('inf')
6        self.counter = 0
7
8    def __call__(self, val_loss):
9        if val_loss < self.best_loss - self.min_delta:
10            self.best_loss = val_loss
11            self.counter = 0
12            return False  # Continue training
13        else:
14            self.counter += 1
15            return self.counter >= self.patience  # Stop if patience exceeded
16
17# Typical result: Training stops at epoch 60-80 instead of 100
18# Time savings: ~20-30%

1import torch.distributed as dist
2from torch.nn.parallel import DistributedDataParallel as DDP
3
4# Initialize distributed training
5dist.init_process_group("nccl")
6model = DDP(model, device_ids=[local_rank])
7
8# Training loop remains similar
9# Data sampler ensures each GPU sees different data

1# Learning rate finder
2from torch_lr_finder import LRFinder
3
4lr_finder = LRFinder(model, optimizer, criterion)
5lr_finder.range_test(train_loader, end_lr=1, num_iter=100)
6lr_finder.plot()  # Find the steepest descent point
7
8# Typical result for AMNL: lr ≈ 1e-3 to 3e-3

1# Save checkpoints efficiently
2def save_checkpoint(model, optimizer, epoch, best_loss, path):
3    torch.save({
4        'epoch': epoch,
5        'model_state_dict': model.state_dict(),
6        'optimizer_state_dict': optimizer.state_dict(),
7        'best_loss': best_loss,
8    }, path)
9
10# Save every N epochs + best model
11if epoch % 10 == 0:
12    save_checkpoint(model, optimizer, epoch, loss, f'checkpoint_{epoch}.pt')
13if loss < best_loss:
14    save_checkpoint(model, optimizer, epoch, loss, 'best_model.pt')