Merge branch 'main' into cws/bxByJz

Author: cwsmith

.github/workflows/checkTrainingParams.yml

@@ -0,0 +1,99 @@
+name: Check Training Parameters
+
+on:
+  pull_request:
+    branches: [ main ]
+    paths:
+      - 'XPointMLTest.py'  # Only trigger when this specific file changes
+
+jobs:
+  check-training-params:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout PR code
+        uses: actions/checkout@v3
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
+          path: pr-code
+          fetch-depth: 1
+
+      - name: Checkout main branch
+        uses: actions/checkout@v3
+        with:
+          ref: main
+          path: main-code
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.10'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install pytest
+
+      - name: Run parameter check
+        run: |
+          python - <<EOF
+import re
+import sys
+
+# Training parameter patterns to check
+params_to_check = {
+    'epochs': r'--epochs\', type=int, default=(\d+),',
+    'trainFrameFirst': r'--trainFrameFirst\', type=int, default=(\d+),',
+    'trainFrameLast': r'--trainFrameLast\', type=int, default=(\d+),',
+    'validationFrameFirst': r'--validationFrameFirst\', type=int, default=(\d+),',
+    'validationFrameLast': r'--validationFrameLast\', type=int, default=(\d+),',
+}
+
+# Files to check
+files_to_check = ['main-code/XPointMLTest.py', 'pr-code/XPointMLTest.py']
+
+main_params = {}
+pr_params = {}
+
+# Extract parameters from main branch
+with open('main-code/XPointMLTest.py', 'r') as f:
+    content = f.read()
+    for param, pattern in params_to_check.items():
+        match = re.search(pattern, content)
+        if match:
+            main_params[param] = int(match.group(1))
+        else:
+            print(f"Warning: Could not find parameter '{param}' in main branch code")
+
+# Extract parameters from PR
+with open('pr-code/XPointMLTest.py', 'r') as f:
+    content = f.read()
+    for param, pattern in params_to_check.items():
+        match = re.search(pattern, content)
+        if match:
+            pr_params[param] = int(match.group(1))
+        else:
+            print(f"Warning: Could not find parameter '{param}' in PR code")
+
+# Compare parameters
+mismatch = False
+for param in params_to_check.keys():
+    if param in main_params and param in pr_params:
+        if main_params[param] != pr_params[param]:
+            print(f"❌ Parameter '{param}' has changed: {main_params[param]} -> {pr_params[param]}")
+            mismatch = True
+        else:
+            print(f"✅ Parameter '{param}' unchanged: {main_params[param]}")
+    else:
+        print(f"⚠️ Could not compare '{param}' - missing from one or both branches")
+
+# Summary
+print("\n=== Parameter Check Summary ===")
+if mismatch:
+    print("❌ Training parameters have been modified!")
+    print("Detected changes to training configuration parameters.")
+    print("Please verify these changes are intentional and approved.")
+    sys.exit(1)
+else:
+    print("✅ All training parameters match the main branch!")
+    sys.exit(0)
+EOF

XPointMLTest.py

@@ -709,6 +709,78 @@ def printCommandLineArgs(args):
         print(f"  {arg}: {getattr(args, arg)}")
     print("}")
 
+# Function to save model checkpoint
+def save_model_checkpoint(model, optimizer, train_loss, val_loss, epoch, checkpoint_dir="checkpoints"):
+    """
+    Save model checkpoint including model state, optimizer state, and training metrics
+    
+    Parameters:
+    model: The neural network model
+    optimizer: The optimizer used for training
+    train_loss: List of training losses
+    val_loss: List of validation losses
+    epoch: Current epoch number
+    checkpoint_dir: Directory to save checkpoints
+    """
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    
+    checkpoint_path = os.path.join(checkpoint_dir, f"xpoint_model_epoch_{epoch}.pt")
+    
+    # Create checkpoint dictionary
+    checkpoint = {
+        'epoch': epoch,
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'train_loss': train_loss,
+        'val_loss': val_loss
+    }
+    
+    # Save checkpoint
+    torch.save(checkpoint, checkpoint_path)
+    print(f"Model checkpoint saved at epoch {epoch} to {checkpoint_path}")
+    
+    # Save the latest model separately for easy loading
+    latest_path = os.path.join(checkpoint_dir, "xpoint_model_latest.pt")
+    torch.save(checkpoint, latest_path)
+    print(f"Latest model saved to {latest_path}")
+
+
+
+# Function to load model checkpoint
+def load_model_checkpoint(model, optimizer, checkpoint_path):
+    """
+    Load model checkpoint
+    
+    Parameters:
+    model: The neural network model to load weights into
+    optimizer: The optimizer to load state into
+    checkpoint_path: Path to the checkpoint file
+    
+    Returns:
+    model: Updated model with loaded weights
+    optimizer: Updated optimizer with loaded state
+    epoch: Last saved epoch number
+    train_loss: List of training losses
+    val_loss: List of validation losses
+    """
+    if not os.path.exists(checkpoint_path):
+        print(f"No checkpoint found at {checkpoint_path}")
+        return model, optimizer, 0, [], []
+    
+    print(f"Loading checkpoint from {checkpoint_path}")
+    checkpoint = torch.load(checkpoint_path)
+    
+    model.load_state_dict(checkpoint['model_state_dict'])
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    
+    epoch = checkpoint['epoch']
+    train_loss = checkpoint['train_loss']
+    val_loss = checkpoint['val_loss']
+    
+    print(f"Loaded checkpoint from epoch {epoch}")
+    return model, optimizer, epoch, train_loss, val_loss
+
+
 def main():
     args = parseCommandLineArgs()
     checkCommandLineArgs(args)
@@ -741,17 +813,35 @@ def main():
     criterion = DiceLoss(smooth=1.0)
     optimizer = optim.Adam(model.parameters(), lr=args.learningRate)
 
+    checkpoint_dir = "checkpoints"
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    latest_checkpoint_path = os.path.join(checkpoint_dir, "xpoint_model_latest.pt")
+    start_epoch = 0
+    train_loss = []
+    val_loss = []
+
+    if os.path.exists(latest_checkpoint_path):
+        model, optimizer, start_epoch, train_loss, val_loss = load_model_checkpoint(
+            model, optimizer, latest_checkpoint_path
+        )
+        print(f"Resuming training from epoch {start_epoch+1}")
+    else:
+        print("Starting training from scratch")
+
     t2 = timer()
     print("time (s) to prepare model: " + str(t2-t1))
 
     train_loss = []
     val_loss = []
-    
+
     num_epochs = args.epochs
-    for epoch in range(num_epochs):
+    for epoch in range(start_epoch, num_epochs):
         train_loss.append(train_one_epoch(model, train_loader, criterion, optimizer, device))
         val_loss.append(validate_one_epoch(model, val_loader, criterion, device))
         print(f"[Epoch {epoch+1}/{num_epochs}]  TrainLoss={train_loss[-1]} ValLoss={val_loss[-1]}")
+        
+        # Save model checkpoint after each epoch
+        save_model_checkpoint(model, optimizer, train_loss, val_loss, epoch+1, checkpoint_dir)
 
     plot_training_history(train_loss, val_loss)
     print("time (s) to train model: " + str(timer()-t2))