oriented-det 0.1.0__py3-none-any.whl

tools/train_example.py

@@ -0,0 +1,244 @@
+"""Example: Training oriented detection models on DOTA dataset.
+
+This script demonstrates how to:
+1. Load DOTA dataset
+2. Set up model and optimizer
+3. Train with the training engine
+4. Evaluate on validation set
+5. Save checkpoints
+"""
+
+import os
+# Set CUDA memory allocation configuration before importing PyTorch
+# This helps with CUDA out of memory errors by using expandable segments
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import argparse
+from pathlib import Path
+import sys
+
+try:
+    import torch
+    import torch.optim as optim
+    from torch.utils.data import DataLoader
+except ImportError as e:
+    print(f"Required dependencies not installed: {e}")
+    print("Please install: pip install torch")
+    sys.exit(1)
+
+from oriented_det import OrientedRCNN, RotatedRetinaNet
+from oriented_det.data import build_dota_loader, DOTADataset
+from oriented_det.train import train, CheckpointManager, MetricTracker
+from oriented_det.train.utils import collate_dota_samples
+
+try:
+    from torch.utils.tensorboard import SummaryWriter
+    TENSORBOARD_AVAILABLE = True
+except ImportError:
+    TENSORBOARD_AVAILABLE = False
+    SummaryWriter = None  # type: ignore
+
+
+def create_model(model_type: str, num_classes: int, backbone: str = "resnet50", pretrained: bool = False):
+    """Create a detection model.
+    
+    Args:
+        model_type: "oriented_rcnn" or "rotated_retinanet"
+        num_classes: Number of detection classes
+        backbone: Backbone name
+        pretrained: Whether to use pretrained backbone
+    """
+    if model_type == "oriented_rcnn":
+        return OrientedRCNN(
+            num_classes=num_classes,
+            backbone_name=backbone,
+            pretrained_backbone=pretrained,
+        )
+    elif model_type == "rotated_retinanet":
+        return RotatedRetinaNet(
+            num_classes=num_classes,
+            backbone_name=backbone,
+            pretrained_backbone=pretrained,
+        )
+    else:
+        raise ValueError(f"Unknown model type: {model_type}")
+
+
+def create_data_loaders(
+    data_root: Path,
+    batch_size: int = 4,
+    num_workers: int = 4,
+    allowed_classes: list = None,
+    difficult_strategy: str = "drop",
+):
+    """Create training and validation data loaders.
+    
+    Args:
+        data_root: Root directory of DOTA dataset
+        batch_size: Batch size
+        num_workers: Number of data loading workers
+        allowed_classes: Optional list of allowed class names
+        difficult_strategy: drop | ignore | keep (DOTA difficult flag handling)
+    """
+    # Note: This is a simplified example. In practice, you'd need to:
+    # 1. Load images properly (DOTADataset returns paths, not tensors)
+    # 2. Apply transforms/augmentation
+    # 3. Create a proper collate function that loads images
+    
+    train_loader = build_dota_loader(
+        root_dir=data_root,
+        split="train",
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        allowed_classes=allowed_classes,
+        difficult_strategy=difficult_strategy,
+    )
+    
+    val_loader = build_dota_loader(
+        root_dir=data_root,
+        split="val",
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        allowed_classes=allowed_classes,
+        difficult_strategy=difficult_strategy,
+    )
+    
+    return train_loader, val_loader
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train oriented detection model on DOTA")
+    parser.add_argument("data_root", type=Path, help="Root directory of DOTA dataset")
+    parser.add_argument("--model-type", choices=["oriented_rcnn", "rotated_retinanet"], default="oriented_rcnn",
+                       help="Model type to train")
+    parser.add_argument("--backbone", default="resnet50", help="Backbone architecture")
+    parser.add_argument("--num-classes", type=int, default=15, help="Number of classes")
+    parser.add_argument("--batch-size", type=int, default=4, help="Batch size")
+    parser.add_argument("--num-epochs", type=int, default=12, help="Number of training epochs")
+    parser.add_argument("--lr", type=float, default=0.001, help="Learning rate")
+    parser.add_argument("--device", default=None,
+                       help="Device to train on")
+    parser.add_argument("--checkpoint-dir", type=Path, default=Path("checkpoints"),
+                       help="Directory to save checkpoints")
+    parser.add_argument("--resume", type=Path, help="Path to checkpoint to resume from")
+    parser.add_argument("--use-amp", action="store_true", help="Use automatic mixed precision")
+    parser.add_argument("--gradient-accumulation", type=int, default=1,
+                       help="Gradient accumulation steps")
+    parser.add_argument("--num-workers", type=int, default=4, help="Number of data loading workers")
+    parser.add_argument("--log-dir", type=Path, help="Directory for TensorBoard logs (enables TensorBoard if provided)")
+    
+    args = parser.parse_args()
+
+    if args.device is None:
+        from oriented_det.utils import get_device
+        args.device = str(get_device())
+
+    # Create model
+    print(f"Creating {args.model_type} model with {args.num_classes} classes...")
+    model = create_model(args.model_type, args.num_classes, args.backbone, pretrained=True)
+    model.to(args.device)
+    
+    # Create data loaders
+    print(f"Loading DOTA dataset from {args.data_root}...")
+    try:
+        train_loader, val_loader = create_data_loaders(
+            args.data_root,
+            batch_size=args.batch_size,
+            num_workers=args.num_workers,
+        )
+        print(f"Train samples: {len(train_loader.dataset)}")
+        print(f"Val samples: {len(val_loader.dataset)}")
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        print("\nNote: This example requires a properly formatted DOTA dataset.")
+        print("The dataset should have the following structure:")
+        print("  data_root/")
+        print("    labelTxt/  (annotation files)")
+        print("    images/    (image files)")
+        sys.exit(1)
+    
+    # Create optimizer
+    optimizer = optim.SGD(
+        model.parameters(),
+        lr=args.lr,
+        momentum=0.9,
+        weight_decay=0.0001,
+    )
+    
+    # Learning rate scheduler
+    lr_scheduler = optim.lr_scheduler.StepLR(
+        optimizer,
+        step_size=8,
+        gamma=0.1,
+    )
+    
+    # Checkpoint manager
+    checkpoint_manager = CheckpointManager(
+        args.checkpoint_dir,
+        best_metric="total_loss",  # In practice, use mAP
+        higher_is_better=False,
+    )
+    
+    # Setup TensorBoard writer if log_dir is provided
+    writer = None
+    if args.log_dir:
+        if not TENSORBOARD_AVAILABLE:
+            print("Warning: TensorBoard is not available. Install with: pip install tensorboard")
+        else:
+            args.log_dir.mkdir(parents=True, exist_ok=True)
+            writer = SummaryWriter(log_dir=str(args.log_dir))
+            print(f"TensorBoard logging enabled. Logs saved to: {args.log_dir}")
+            print(f"View logs with: tensorboard --logdir {args.log_dir}")
+    
+    # Resume from checkpoint if provided
+    start_epoch = 0
+    if args.resume and args.resume.exists():
+        print(f"Resuming from checkpoint: {args.resume}")
+        checkpoint_manager.load(args.resume, model, optimizer)
+        start_epoch = checkpoint_manager.load(args.resume, model, optimizer).get("epoch", 0) + 1
+    
+    # Training configuration
+    print("\nTraining configuration:")
+    print(f"  Model: {args.model_type}")
+    print(f"  Backbone: {args.backbone}")
+    print(f"  Batch size: {args.batch_size}")
+    print(f"  Learning rate: {args.lr}")
+    print(f"  Epochs: {args.num_epochs}")
+    print(f"  Device: {args.device}")
+    print(f"  Mixed precision: {args.use_amp}")
+    print(f"  Gradient accumulation: {args.gradient_accumulation}")
+    print()
+    
+    # Train
+    try:
+        history = train(
+            model=model,
+            train_loader=train_loader,
+            optimizer=optimizer,
+            device=torch.device(args.device),
+            num_epochs=args.num_epochs,
+            val_loader=val_loader,
+            lr_scheduler=lr_scheduler,
+            checkpoint_manager=checkpoint_manager,
+            use_amp=args.use_amp,
+            gradient_accumulation_steps=args.gradient_accumulation,
+            max_grad_norm=1.0,
+            start_epoch=start_epoch,
+            writer=writer,
+        )
+        
+        print("\nTraining completed!")
+        print(f"Checkpoints saved to: {args.checkpoint_dir}")
+        
+    except KeyboardInterrupt:
+        print("\nTraining interrupted by user")
+        print(f"Checkpoints saved to: {args.checkpoint_dir}")
+    except Exception as e:
+        print(f"\nTraining error: {e}")
+        raise
+
+
+if __name__ == "__main__":
+    main()

tools/train_multi_gpu.py

@@ -0,0 +1,367 @@
+#!/usr/bin/env python3
+"""Multi-GPU training wrapper for train.py.
+
+This script launches train.py with torchrun for distributed training across multiple GPUs.
+Configuration is read from JSON config files; pass --config <path> and other train.py arguments.
+
+Usage:
+    python train_multi_gpu.py [--nproc-per-node N] --config <config.json> [train.py arguments...]
+
+Examples:
+    # Use all available GPUs with default config
+    python train_multi_gpu.py --config configs/rotated_faster_rcnn/dota_le90_3x.json
+
+    # Use 4 GPUs
+    python train_multi_gpu.py --nproc-per-node 4 --config configs/rotated_faster_rcnn/dota_le90_3x.json
+
+    # Use 2 GPUs with custom arguments
+    python train_multi_gpu.py --nproc-per-node 2 --config configs/.../config.json --use-amp
+
+Checkpoint loading is configured only in the JSON ``checkpoint`` section (same as ``train.py``).
+
+The script automatically:
+- Detects the number of available GPUs
+- Launches train.py with torchrun
+- Forwards all arguments (including --config) to train.py
+- Sets up proper environment variables for distributed training
+"""
+
+import subprocess
+import sys
+import os
+from pathlib import Path
+
+try:
+    import torch
+except ImportError:
+    print("Error: PyTorch is not installed. Please install PyTorch first.")
+    sys.exit(1)
+
+
+def get_num_gpus():
+    """Get the number of available GPUs."""
+    if not torch.cuda.is_available():
+        print("Warning: CUDA is not available. Falling back to CPU training.")
+        return 0
+    return torch.cuda.device_count()
+
+
+def find_torchrun():
+    """Find torchrun executable."""
+    # Try torchrun first (PyTorch 1.9+)
+    try:
+        result = subprocess.run(
+            ["torchrun", "--help"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+        if result.returncode == 0:
+            return "torchrun"
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        pass
+    
+    # Fallback to python -m torch.distributed.launch
+    try:
+        result = subprocess.run(
+            [sys.executable, "-m", "torch.distributed.launch", "--help"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+        if result.returncode == 0:
+            return "python -m torch.distributed.launch"
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        pass
+    
+    return None
+
+
+def main():
+    """Main entry point."""
+    import argparse
+    
+    parser = argparse.ArgumentParser(
+        description="Multi-GPU training wrapper for train.py (config-based training)",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Use all available GPUs (pass --config and any train.py options)
+  python train_multi_gpu.py --config configs/rotated_faster_rcnn/dota_le90_3x.json
+
+  # Use 4 GPUs
+  python train_multi_gpu.py --nproc-per-node 4 --config configs/.../config.json
+
+  # Use 2 GPUs with custom arguments
+  python train_multi_gpu.py --nproc-per-node 2 --config configs/.../config.json --use-amp --debug
+        """,
+    )
+    
+    parser.add_argument(
+        "--nproc-per-node",
+        type=int,
+        default=None,
+        help="Number of processes per node (default: number of available GPUs)",
+    )
+    
+    parser.add_argument(
+        "--master-port",
+        type=int,
+        default=29500,
+        help="Master port for distributed training (default: 29500)",
+    )
+    
+    parser.add_argument(
+        "--backend",
+        choices=("gloo", "nccl"),
+        default=None,
+        help="Distributed backend (default: nccl for single-node, faster than Gloo with P2P/SHM enabled)",
+    )
+    
+    # Parse known args to separate wrapper args from train.py args
+    args, unknown_args = parser.parse_known_args()
+    
+    # Get number of GPUs
+    num_gpus = get_num_gpus()
+    
+    if num_gpus == 0:
+        print("Error: No GPUs available. Cannot run multi-GPU training.")
+        print("Please use train.py directly for CPU or single-GPU training.")
+        sys.exit(1)
+    
+    # Determine number of processes
+    nproc_per_node = args.nproc_per_node or num_gpus
+    
+    if nproc_per_node > num_gpus:
+        print(f"Warning: Requested {nproc_per_node} GPUs but only {num_gpus} available.")
+        print(f"Using {num_gpus} GPUs instead.")
+        nproc_per_node = num_gpus
+    
+    if nproc_per_node < 1:
+        print("Error: nproc-per-node must be at least 1.")
+        sys.exit(1)
+    
+    # Find torchrun
+    torchrun_cmd = find_torchrun()
+    if torchrun_cmd is None:
+        print("Error: Could not find torchrun or torch.distributed.launch.")
+        print("Please ensure PyTorch is properly installed.")
+        sys.exit(1)
+    
+    # Train via installable module (odet train / oriented_det.cli.train)
+    train_module = "oriented_det.cli.train"
+    extra_args = unknown_args
+    if torchrun_cmd == "torchrun":
+        cmd = [
+            "torchrun",
+            "--nproc-per-node", str(nproc_per_node),
+            "--master-port", str(args.master_port),
+            "-m", train_module,
+        ] + extra_args
+    else:
+        cmd = [
+            sys.executable,
+            "-m", "torch.distributed.launch",
+            "--nproc-per-node", str(nproc_per_node),
+            "--master-port", str(args.master_port),
+            train_module,
+        ] + extra_args
+    
+    # Set up environment for single-node multi-GPU training
+    env = os.environ.copy()
+    
+    # Short TMPDIR to avoid Gloo "File name too long" (store paths can embed hostname)
+    if "TMPDIR" not in env:
+        for d in ("/dev/shm", "/tmp"):
+            if os.path.isdir(d):
+                env["TMPDIR"] = d
+                break
+    
+    # Backend: default NCCL for all (faster with P2P/SHM enabled on single-node A100)
+    # With PyTorch 2.3+ and proper configuration, NCCL works reliably on single-node
+    if args.backend:
+        env["DIST_BACKEND"] = args.backend
+    else:
+        env.setdefault("DIST_BACKEND", "nccl")  # NCCL is faster than Gloo on single-node multi-GPU
+    
+    # Ensure MASTER_ADDR and MASTER_PORT are set for distributed init
+    env.setdefault("MASTER_ADDR", "127.0.0.1")
+    env.setdefault("MASTER_PORT", str(args.master_port))
+    
+    socket_ifname = None
+    use_gib = False
+
+    # Configure NCCL for GCP. gIB is only available on A3 Ultra, A4, A4X instances.
+    # A2 instances (A100) do NOT support gIB - use socket-only NCCL (NVLink intra-node, sockets/gVNIC inter-node).
+    if env.get("DIST_BACKEND", "nccl") == "nccl":
+        # Check instance type to determine if gIB is available
+        try:
+            instance_type = subprocess.run(
+                ["curl", "-s", "-H", "Metadata-Flavor: Google",
+                 "http://metadata.google.internal/computeMetadata/v1/instance/machine-type"],
+                capture_output=True, text=True, timeout=2
+            )
+            if instance_type.returncode == 0:
+                instance_type = instance_type.stdout.strip().split("/")[-1]
+                # gIB is only available on A3 Ultra, A4, A4X (not A2)
+                if instance_type.startswith(("a3-ultra", "a4", "a4x")):
+                    gib_script = "/usr/local/gib/scripts/set_nccl_env.sh"
+                    if os.path.exists(gib_script):
+                        use_gib = True
+                        env["NCCL_NET"] = "gIB"
+                        env["NCCL_CROSS_NIC"] = "0"
+                        env["NCCL_NET_GDR_LEVEL"] = "PIX"
+                        env["NCCL_P2P_NET_CHUNKSIZE"] = "131072"
+                        env["NCCL_NVLS_CHUNKSIZE"] = "524288"
+                        env["NCCL_IB_ADAPTIVE_ROUTING"] = "1"
+                        env["NCCL_IB_QPS_PER_CONNECTION"] = "4"
+                        env["NCCL_IB_TC"] = "52"
+                        env["NCCL_IB_FIFO_TC"] = "84"
+                        env["NCCL_TUNER_CONFIG_PATH"] = "/usr/local/gib/configs/tuner_config_a3u.txtpb"
+                        print(f"Using GCP gIB backend (instance: {instance_type})")
+                    else:
+                        print(f"Warning: gIB script not found on {instance_type}, using socket-only NCCL")
+                else:
+                    print(f"Instance {instance_type} does not support gIB (A2/A100), using socket-only NCCL")
+            else:
+                print("Could not detect instance type, using socket-only NCCL")
+        except Exception:
+            print("Could not detect instance type, using socket-only NCCL")
+
+        # A2 / socket-only: recommended NCCL settings (GDR off, socket tuning, NCCL 2.12.7+)
+        if not use_gib:
+            env.pop("NCCL_NET", None)  # Ensure no inherited gIB on A2
+            env["NCCL_NET_GDR_LEVEL"] = "0"  # Disable GPUDirect RDMA (unsupported on A2)
+            env["NCCL_SOCKET_NTHREADS"] = "4"
+            env["NCCL_NSOCKETS_PER_PEER"] = "4"
+            env.setdefault("NCCL_DEBUG", "WARN")  # WARN: minimal output; use INFO for debugging
+            
+            # Remove GCP gIB shim from LD_LIBRARY_PATH on A2 (shim enforces PIX, incompatible with A2)
+            # The shim (libnccl-net_internal.so) checks guest_config.txtpb and enforces gIB settings
+            ld_path = env.get("LD_LIBRARY_PATH", "")
+            if ld_path:
+                # Remove /usr/local/gib/lib64 to prevent shim from loading
+                paths = [p for p in ld_path.split(":") if p and not p.rstrip("/").endswith("/usr/local/gib/lib64")]
+                if paths:
+                    env["LD_LIBRARY_PATH"] = ":".join(paths)
+                else:
+                    env.pop("LD_LIBRARY_PATH", None)
+            print("Removed GCP gIB shim from LD_LIBRARY_PATH (A2 uses socket-only NCCL)")
+
+        # Detect and set network interface for single-node communication (required for socket-only NCCL)
+        socket_ifname = env.get("NCCL_SOCKET_IFNAME")
+        if not socket_ifname:
+            try:
+                result = subprocess.run(["ip", "route", "get", "8.8.8.8"], capture_output=True, text=True, timeout=2)
+                if result.returncode == 0:
+                    for line in result.stdout.split('\n'):
+                        if 'dev' in line:
+                            parts = line.split()
+                            if 'dev' in parts:
+                                idx = parts.index('dev')
+                                if idx + 1 < len(parts):
+                                    socket_ifname = parts[idx + 1]
+                                    break
+            except (FileNotFoundError, subprocess.TimeoutExpired, Exception):
+                pass
+        if not socket_ifname:
+            for ifname in ["ens8", "ens4", "eth0", "eno1"]:
+                try:
+                    with open(f"/sys/class/net/{ifname}/operstate", "r") as f:
+                        if "up" in f.read().lower():
+                            socket_ifname = ifname
+                            break
+                except (FileNotFoundError, IOError):
+                    continue
+        if socket_ifname:
+            env["NCCL_SOCKET_IFNAME"] = socket_ifname
+    
+    # For Gloo backend: set GLOO_SOCKET_IFNAME (reuse socket_ifname if already detected for NCCL)
+    if env.get("DIST_BACKEND", "nccl") == "gloo":
+        if not socket_ifname:
+            try:
+                result = subprocess.run(["ip", "route", "get", "8.8.8.8"], capture_output=True, text=True, timeout=2)
+                if result.returncode == 0:
+                    for line in result.stdout.split('\n'):
+                        if 'dev' in line:
+                            parts = line.split()
+                            if 'dev' in parts:
+                                idx = parts.index('dev')
+                                if idx + 1 < len(parts):
+                                    socket_ifname = parts[idx + 1]
+                                    break
+            except (FileNotFoundError, subprocess.TimeoutExpired, Exception):
+                pass
+            if not socket_ifname:
+                for ifname in ["ens8", "ens4", "eth0", "eno1"]:
+                    try:
+                        with open(f"/sys/class/net/{ifname}/operstate", "r") as f:
+                            if "up" in f.read().lower():
+                                socket_ifname = ifname
+                                break
+                    except (FileNotFoundError, IOError):
+                        continue
+    
+    # Optimize Gloo backend settings for better performance
+    if env.get("DIST_BACKEND", "nccl") == "gloo":
+        # Use proper network interface for Gloo
+        if socket_ifname:
+            env.setdefault("GLOO_SOCKET_IFNAME", socket_ifname)
+        # Increase Gloo timeout (default 30s, increase for large models)
+        env.setdefault("GLOO_TIMEOUT_SECONDS", "1800")  # 30 minutes
+        # Use async operations where possible (if supported)
+        # Note: GLOO_ASYNC may not be available in all PyTorch versions
+    
+    # Print info
+    print("=" * 80)
+    print("Multi-GPU Training Launcher")
+    print("=" * 80)
+    print(f"GPUs available: {num_gpus}")
+    print(f"GPUs to use: {nproc_per_node}")
+    print(f"Module: {train_module}")
+    proj = env.get("ORIENTED_DET_PROJECT_ROOT")
+    if proj:
+        print(f"ORIENTED_DET_PROJECT_ROOT: {proj}")
+    from oriented_det.pretrained import get_pretrained_dir
+
+    pre = env.get("ORIENTED_DET_PRETRAINED_DIR")
+    print(f"Pretrained cache: {pre or get_pretrained_dir()}")
+    print(f"Command: {' '.join(cmd)}")
+    print("=" * 80)
+    print("Distributed Training Configuration:")
+    backend = env.get('DIST_BACKEND', 'nccl')
+    print(f"  Backend: {backend} (default: NCCL, use --backend gloo for Gloo)")
+    print(f"  MASTER_ADDR: {env.get('MASTER_ADDR', '127.0.0.1')}")
+    print(f"  MASTER_PORT: {env.get('MASTER_PORT', '29500')}")
+    if backend == "nccl":
+        net = env.get("NCCL_NET", "socket (default)")
+        ifname = env.get("NCCL_SOCKET_IFNAME", "not set")
+        if net == "gIB":
+            gdr_level = env.get("NCCL_NET_GDR_LEVEL", "not set")
+            print(f"  NCCL_NET: {net} (gIB backend)")
+            print(f"  NCCL_NET_GDR_LEVEL: {gdr_level}")
+        else:
+            print(f"  NCCL_NET: {net} (socket-only, A2)")
+            print(f"  NCCL_NET_GDR_LEVEL: {env.get('NCCL_NET_GDR_LEVEL', 'not set')}")
+            print(f"  NCCL_SOCKET_NTHREADS: {env.get('NCCL_SOCKET_NTHREADS', 'not set')}")
+            print(f"  NCCL_NSOCKETS_PER_PEER: {env.get('NCCL_NSOCKETS_PER_PEER', 'not set')}")
+            print(f"  NCCL_DEBUG: {env.get('NCCL_DEBUG', 'not set')}")
+        print(f"  NCCL_SOCKET_IFNAME: {ifname}")
+    print("=" * 80)
+    print()
+    
+    # Launch training
+    try:
+        # Run the command with environment variables and forward stdout/stderr
+        result = subprocess.run(cmd, env=env, check=False)
+        sys.exit(result.returncode)
+    except KeyboardInterrupt:
+        print("\nTraining interrupted by user.")
+        sys.exit(130)
+    except Exception as e:
+        print(f"Error launching training: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()