File size: 5,082 Bytes
00801a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | #!/usr/bin/env python3
"""
ACL-LKNet Training CLI
======================
Standalone command-line entrypoint for training ACL-LKNet models.
Supports single-fold training, full 5-fold stratified cross-validation,
and Phase 1 self-supervised pretraining (Masked Slice Modeling).
Usage Examples:
# Train Fold 1 with default configuration:
python train_cv.py --data_dir /path/to/mrnet --fold 1
# Train all 5 folds sequentially for cross-validation:
python train_cv.py --data_dir /path/to/mrnet --cv
# Run Phase 1 Masked Slice Modeling (SSL) pretraining:
python train_cv.py --data_dir /path/to/mrnet --ssl_pretrain
"""
import os
import sys
import argparse
import logging
import torch
# Ensure local package imports work seamlessly
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from src.config import Config
from src.train import train_supervised, train_ssl, train_5fold_cross_validation
from src.utils import set_seed, setup_logging
def parse_args():
parser = argparse.ArgumentParser(
description="Train ACL-LKNet (Phase 1 SSL Pretraining or Phase 2 Supervised Fine-Tuning)."
)
parser.add_argument(
"--data_dir", type=str, default="./data/mrnet",
help="Path to Stanford MRNet dataset root directory (containing train, valid, and test folders)."
)
parser.add_argument(
"--output_dir", type=str, default="./checkpoints",
help="Directory where model checkpoints and logs will be saved."
)
parser.add_argument(
"--backbone", type=str, default="convnext_tiny",
choices=["convnext_tiny", "resnet18", "replknet"],
help="Backbone architecture (ConvNeXt-Tiny recommended for large receptive field)."
)
parser.add_argument(
"--fold", type=int, default=1, choices=[1, 2, 3, 4, 5],
help="Fold index to train (1-5) when not running full cross-validation."
)
parser.add_argument(
"--cv", action="store_true",
help="Run complete 5-fold stratified cross-validation."
)
parser.add_argument(
"--ssl_pretrain", action="store_true",
help="Run Phase 1 Masked Slice Modeling (MSM) self-supervised pretraining."
)
parser.add_argument(
"--ssl_checkpoint", type=str, default=None,
help="Optional path to SSL-pretrained backbone checkpoint to initialize Phase 2 training."
)
parser.add_argument(
"--epochs", type=int, default=30,
help="Maximum training epochs (monitored by early stopping patience=20)."
)
parser.add_argument(
"--lr_backbone", type=float, default=1.0e-5,
help="Differential learning rate for pretrained backbone."
)
parser.add_argument(
"--lr_head", type=float, default=3.0e-4,
help="Learning rate for newly initialized attention and classification heads."
)
parser.add_argument(
"--batch_size", type=int, default=1,
help="Physical batch size (keep = 1 on 15-16 GB GPUs to prevent OOM)."
)
parser.add_argument(
"--accum_steps", type=int, default=8,
help="Gradient accumulation steps to achieve effective batch size = 8."
)
parser.add_argument(
"--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
help="Compute device ('cuda' or 'cpu')."
)
parser.add_argument(
"--seed", type=int, default=42,
help="Random seed for reproducibility across PyTorch, NumPy, and Python."
)
return parser.parse_args()
def main():
args = parse_args()
# Initialize master configuration
config = Config(
data_dir=args.data_dir,
checkpoint_dir=args.output_dir,
log_dir=os.path.join(args.output_dir, "logs"),
backbone_name=args.backbone,
epochs=args.epochs,
backbone_lr=args.lr_backbone,
lr=args.lr_head,
batch_size=args.batch_size,
accumulation_steps=args.accum_steps,
device=args.device,
seed=args.seed,
)
os.makedirs(config.checkpoint_dir, exist_ok=True)
os.makedirs(config.log_dir, exist_ok=True)
setup_logging(config.log_dir)
set_seed(config.seed)
logging.info(f"Initialized ACL-LKNet with Backbone: {config.backbone_name} on {config.device}")
if args.ssl_pretrain:
logging.info("Starting Phase 1: Masked Slice Modeling (MSM) SSL Pretraining...")
train_ssl(config)
elif args.cv:
logging.info(f"Starting 5-Fold Stratified Cross-Validation on MRNet...")
results = train_5fold_cross_validation(config, ssl_checkpoint=args.ssl_checkpoint)
print("\n" + "=" * 60)
print("Cross-Validation Complete! Summary:")
print(f"Mean Val AUROC: {results.get('mean_val_auc', 'N/A')}")
print("=" * 60)
else:
logging.info(f"Starting Single Fold Training: Fold {args.fold}...")
config.experiment_name = f"acl_lknet_fold{args.fold}"
train_supervised(config, fold_idx=args.fold - 1, ssl_checkpoint=args.ssl_checkpoint)
if __name__ == "__main__":
main()
|