pytorch-transfer-learning-debugger 0.1.1

A debugger for PyTorch transfer-learning & fine-tuning jobs.

PyTorch Transfer Learning Debugger (v0.1.1)

Diagnose and debug common issues when performing in transfer learning & fine-tuning with PyTorch.

An interactive debugging experience for PyTorch code, enabling users to diagnose issues related to weight updates during training, as well as to address problems with slow or unstable convergence.

---

Installation (PyPI)

pip install the package:

pip install pytorch-transfer-learning-debugger

or

pip install git+https://github.com/ZachWolpe/PyTorch-transfer-learning-debugger.git

Access the PyPI source.

---

Benefits

• Helps identify if weights are properly updating during training
• Provides visibility into each stage of the training process
• Enables early detection of training issues
• Supports both high-level and granular debugging approaches

This debugging tool is particularly useful when:

• Troubleshooting training convergence issues
• Verifying proper weight updates
• Understanding the flow of data through the model
• Diagnosing optimization problems

---

Getting Started

Purpose

An interactive debugging tool for PyTorch code that helps diagnose issues related to weight updates during training and addresses problems with slow or unstable convergence.

Key Components

1. Debug Mode Features

The debugger is implemented through a torch_debugger class that provides two main debugging modes:

1. Basic Debug Mode (debug_mode=True):

   • Tracks model weight updates
   • Only monitors weights with requires_grad=True
   • Populates torch_debugger._track_weights dictionary for analysis

2. Granular Debug Mode (granular_logging=True):

   • Exits after one cycle through the dataloader
   • Logs each stage of the computational graph:
     • Optimizer verification
     • Gradient zeroing
     • Data movement to device
     • Forward pass and loss computation
     • Gradient computation and optimizer steps

2. Implementation Example

Here's how the debugger is integrated into the training loop:

def train_model(model, criterion, optimizer, scheduler, num_epochs=25, 
                debug_mode=True, granular_logging=False):
    # Initialize debugger if debug_mode is enabled
    torch_debugger_inst = None
    if debug_mode:
        torch_debugger_inst = torch_debugger(enable_log_stage=granular_logging)
        torch_debugger_inst.verify_optimizer_state(optimizer)
        torch_debugger_inst.initial_model_weights_state = model.state_dict().copy()
        
    # Training loop with debugging hooks
    for epoch in range(num_epochs):
        for phase in ['train', 'val']:
            for inputs, labels in dataloaders[phase]:
                if debug_mode:
                    torch_debugger_inst.log_stage('Moving data to device')
                    # ... more logging stages ...
                    
                if phase == 'train':
                    torch_debugger_inst.track_weights(model.named_parameters())

3. Usage

To use the debugger:

torch_debugger_inst, model = train_model(
    model_ft,
    criterion,
    optimizer_ft,
    exp_lr_scheduler,
    num_epochs=25,
    debug_mode=True,
    granular_logging=True  # For detailed stage logging
)

4. Analysis Tools

The debugger provides visualization tools to analyze training:

# Plot weight updates over time
torch_debugger_inst.plot_weight_updates()

# Access tracked weights
weights = torch_debugger_inst._track_weights['conv1.weight']

---

Comprehensive Tutorial

For a comprehensive tutorial, see the GettingStarted.ipynb Juypyter notebook.

---

Potential Issues during Transfer Learning

Category	Issue	Symptoms	Solutions
Learning Rate	Improper Configuration	• Too high: unstable training • Too low: slow learning/stuck • Poor adjustment for transfer	• Start with 10-3 or 10-4 of original LR • Use LR finder • Implement LR scheduling • Use different LRs for layers
Layer Management	Freezing Issues	• Wrong layers frozen • Too many/few layers frozen • No gradual unfreezing • Model not learning	• Start by freezing all but final layers • Gradually unfreeze from top • Monitor layer gradients • Selective layer unfreezing
Data Preparation	Input Problems	• Runtime errors • Poor model performance • Slow convergence • Class imbalance	• Adjust transforms • Match source preprocessing • Balance classes • Use validation set • Data augmentation
Architecture	Structure Mismatches	• Bad final layer modifications • Size/channel mismatches • Poor layer initialization • Training fails	• Verify input/output dimensions • Use proper initialization • Match pretrained architecture • Add adaptation layers
Training Process	Optimization Issues	• Not converging • Unstable training • Poor generalization • Memory errors	• Use Adam/AdamW for fine-tuning • Verify loss matches task • Start with small batches • Use gradient accumulation
Implementation	Technical Errors	• Model not in train mode • Gradients not zeroed • Wrong device (CPU/GPU) • Memory leaks	• Use training checklist • Implement proper train/eval • Check device placement • Monitor memory usage
Resource Management	Performance Issues	• OOM errors • Slow training • Resource inefficiency	• Batch size adjustment • Gradient checkpointing • Efficient data loading • Use GPU acceleration
Monitoring & Validation	Quality Control	• Poor metric tracking • Missing validation • Overfitting • Inconsistent results	• Use debugging tools • Implement validation loops • Track multiple metrics • Monitor gradient flow
Pretrained Model	Weight Issues	• Wrong pretrained weights • Corrupted weights • Version incompatibility	• Verify model source • Check model checksums • Match framework versions
Feature Extraction	Transfer Problems	• Suboptimal transfer • Poor adaptation • Loss of pretrained features	• Choose appropriate layers • Add adaptation layers • Fine-tune feature extractors • Knowledge distillation

---

Potential Issues with the Computational Graph

Issue Category	Specific Problem	Symptoms	Solutions
Gradient Flow	Vanishing Gradients	- Near-zero gradients in early layers - Model not learning	- Use gradient clipping - Add residual connections - Change activation functions
	Exploding Gradients	- NaN losses - Large gradient values	- Implement gradient clipping - Reduce learning rate - Check initialization
	Disconnected Graphs	- Some parameters not updating - Partial learning	- Remove accidental detach() - Verify tensor operations maintain graph
Tensor Operations	In-place Operations	- Backward pass errors - "Leaf variable modified" error	- Replace in-place ops (+=) with regular ops (+) - Create new tensors instead of modifying
	Device Mismatches	- Runtime errors - CUDA errors	- Use .to(device) consistently - Implement device checker
	Detached Tensors	- No gradients flowing - Parts of model not learning	- Set requires_grad=True where needed - Remove unnecessary detach() calls
Autograd Engine	Broken Computational Paths	- Gradients not computed - backward() errors	- Fix graph connections - Ensure proper tensor operations
	Mixed Precision Errors	- NaN losses - Unstable training	- Use GradScaler - Adjust loss scaling - Check dtype consistency
Loss Computation	Zero/NaN Losses	- Model not learning - Training instability	- Verify loss function implementation - Check input normalization
	Wrong Reduction	- Incorrect gradient scaling - Slow convergence	- Set appropriate reduction method - Verify batch dimension handling

---

: [e-mail]
: 09.12.2024