flashstudio 0.1.0__py3-none-any.whl

flashstudio/pages/model.py

@@ -0,0 +1,370 @@
+"""FlashStudio — Model Selection & Finetune Configuration Page."""
+
+import streamlit as st
+
+
+FLASHDET_MODELS = {
+    "FlashDet-Pico": {
+        "size": "p", "params": "~298K", "speed": "Ultra-fast",
+        "backbone": "LiteBackbone (0.5x)", "neck": "PicoNeck (64ch)",
+        "head": "E2EDualHead", "best_for": "Edge / Mobile / MCU",
+    },
+    "FlashDet-Nano": {
+        "size": "n", "params": "~790K", "speed": "Very fast",
+        "backbone": "FlashBackbone (stem=32)", "neck": "PicoNeck (96ch)",
+        "head": "E2EDualHead", "best_for": "Embedded / IoT",
+    },
+    "FlashDet-Small": {
+        "size": "s", "params": "~1.8M", "speed": "Fast",
+        "backbone": "FlashBackbone (stem=48)", "neck": "PicoNeck (128ch)",
+        "head": "E2EDualHead", "best_for": "General purpose",
+    },
+    "FlashDet-Medium": {
+        "size": "m", "params": "~3.6M", "speed": "Balanced",
+        "backbone": "FlashBackbone (stem=64)", "neck": "PicoNeck (192ch)",
+        "head": "E2EDualHead", "best_for": "High accuracy",
+    },
+    "FlashDet-Large": {
+        "size": "l", "params": "~5.8M", "speed": "Accurate",
+        "backbone": "FlashBackbone (stem=80)", "neck": "PicoNeck (256ch)",
+        "head": "E2EDualHead", "best_for": "High accuracy",
+    },
+    "FlashDet-X": {
+        "size": "x", "params": "~9.0M", "speed": "Max accuracy",
+        "backbone": "FlashBackbone (stem=96)", "neck": "PicoNeck (320ch)",
+        "head": "E2EDualHead", "best_for": "Max accuracy / Server",
+    },
+}
+
+OTHER_ARCHITECTURES = {
+    "YOLOv8": {"params": "Varies", "speed": "Fast", "best_for": "General YOLO"},
+    "YOLOv9": {"params": "Varies", "speed": "Fast", "best_for": "PGI-based detection"},
+    "YOLOv10": {"params": "Varies", "speed": "Fast", "best_for": "PSA-enhanced"},
+    "YOLOv11": {"params": "Varies", "speed": "Fast", "best_for": "C2PSA-based"},
+    "YOLOX": {"params": "Varies", "speed": "Fast", "best_for": "Anchor-free YOLO"},
+}
+
+OPTIMIZERS = ["AdamW", "SGD", "MuSGD"]
+LORA_VARIANTS = ["standard", "dora", "lora_plus", "adalora", "ortho", "lora_fa"]
+
+
+def render_model_page():
+    """Render model selection and fine-tuning configuration."""
+    from flashstudio.components.styles import render_page_header
+    render_page_header("🧠", "Model Configuration",
+                       "Choose your architecture and configure all training parameters.")
+
+    tab_arch, tab_hyper, tab_aug, tab_advanced = st.tabs([
+        "🏗️ Architecture", "⚙️ Hyperparameters", "🎨 Augmentation", "🔧 Advanced"
+    ])
+
+    with tab_arch:
+        _render_architecture_tab()
+
+    with tab_hyper:
+        _render_hyperparameters_tab()
+
+    with tab_aug:
+        _render_augmentation_tab()
+
+    with tab_advanced:
+        _render_advanced_tab()
+
+    st.divider()
+    _render_config_summary()
+
+
+def _render_architecture_tab():
+    """Model architecture selection."""
+    st.markdown("### Select Architecture")
+
+    arch_family = st.radio(
+        "Architecture Family",
+        ["FlashDet (recommended)", "YOLOv8", "YOLOv9", "YOLOv10", "YOLOv11", "YOLOX"],
+        key="arch_family",
+        horizontal=True,
+    )
+
+    if arch_family == "FlashDet (recommended)":
+        _render_flashdet_config()
+    else:
+        _render_yolo_config(arch_family.replace(" (recommended)", ""))
+
+
+def _render_flashdet_config():
+    """FlashDet-specific architecture configuration."""
+    col_select, col_info = st.columns([1, 2])
+
+    with col_select:
+        model_size = st.radio(
+            "Model Size",
+            list(FLASHDET_MODELS.keys()),
+            index=0,
+            key="model_arch",
+        )
+
+        info = FLASHDET_MODELS[model_size]
+
+        # Only Pico has backbone choice
+        if model_size == "FlashDet-Pico":
+            st.selectbox(
+                "Backbone",
+                ["LiteBackbone (ShuffleNetV2-0.5x)", "PicoBackbone (RepNeXt-style)"],
+                key="pico_backbone",
+                help="Only FlashDet-Pico allows backbone selection",
+            )
+        else:
+            st.info(f"**Backbone:** {info['backbone']} (fixed for {model_size})")
+
+        # Neck and Head are ALWAYS fixed — show as info, not selectable
+        st.info(f"**Neck:** {info['neck']} (fixed)")
+        st.info(f"**Head:** {info['head']} (fixed)")
+
+    with col_info:
+        st.markdown(f"### {model_size}")
+
+        metric_cols = st.columns(3)
+        with metric_cols[0]:
+            st.metric("Parameters", info["params"])
+        with metric_cols[1]:
+            st.metric("Speed", info["speed"])
+        with metric_cols[2]:
+            st.metric("Best For", info["best_for"])
+
+        with st.container(border=True):
+            st.markdown("**Architecture Details**")
+            st.caption(f"Backbone: {info['backbone']}")
+            st.caption(f"Neck: {info['neck']}")
+            st.caption(f"Head: {info['head']} (o2o + o2m dual-head, DFL-free)")
+            st.caption("Loss: CIoU + BCE + L1 with STAL assignment")
+            st.caption("Training: ProgLoss (progressive o2m→o2o balancing)")
+
+        _render_pretrain_and_finetune()
+
+
+def _render_yolo_config(arch_name):
+    """YOLO-family architecture configuration."""
+    col_select, col_info = st.columns([1, 2])
+
+    with col_select:
+        st.markdown(f"#### {arch_name} Configuration")
+
+        st.slider("Width Multiplier", 0.25, 1.5, 0.5, step=0.25, key="yolo_width_mult",
+                  help="Controls channel width (0.25=nano, 0.5=small, 1.0=medium)")
+        st.slider("Depth Multiplier", 0.33, 1.5, 0.33, step=0.33, key="yolo_depth_mult",
+                  help="Controls network depth (0.33=small, 0.67=medium, 1.0=large)")
+
+        if arch_name == "YOLOv9":
+            st.checkbox("Enable PGI (Programmable Gradient Info)", value=True, key="yolo_use_pgi")
+        elif arch_name == "YOLOv10":
+            st.checkbox("Enable PSA (Partial Self-Attention)", value=True, key="yolo_use_psa")
+        elif arch_name == "YOLOv11":
+            st.checkbox("Enable C2PSA", value=True, key="yolo_use_c2psa")
+
+    with col_info:
+        yolo_info = OTHER_ARCHITECTURES[arch_name]
+        st.markdown(f"### {arch_name}")
+        st.caption(f"Best for: {yolo_info['best_for']}")
+        st.caption("Backbone, Neck, Head are fixed per architecture — not user-selectable.")
+
+        with st.container(border=True):
+            st.markdown("**Note**")
+            st.write(
+                f"The {arch_name} architecture uses its own fixed backbone, neck, and head. "
+                "You can control model capacity via width/depth multipliers."
+            )
+
+        _render_pretrain_and_finetune()
+
+
+def _render_pretrain_and_finetune():
+    """Shared pretrained weights and fine-tuning strategy UI."""
+    with st.container(border=True):
+        st.markdown("**Pretrained Weights**")
+        pretrain_opt = st.radio(
+            "Initialize from",
+            ["COCO pretrained (recommended)", "ImageNet backbone only",
+             "Random (train from scratch)", "Custom weights"],
+            key="pretrain_option",
+        )
+        if pretrain_opt == "Custom weights":
+            st.text_input("Path to weights", placeholder="/content/model.pth", key="custom_weights")
+
+    with st.container(border=True):
+        st.markdown("**Fine-tuning Strategy**")
+        st.radio(
+            "Training mode",
+            [
+                "Full fine-tune (all layers trainable)",
+                "Freeze backbone (train neck + head only)",
+                "Freeze backbone + neck (train head only)",
+                "LoRA fine-tune (low-rank adapters)",
+            ],
+            key="finetune_strategy",
+        )
+        if st.session_state.get("finetune_strategy", "").startswith("LoRA"):
+            st.selectbox("LoRA Variant", LORA_VARIANTS, key="lora_variant")
+            st.slider("LoRA Rank", 4, 64, 16, key="lora_rank")
+            st.slider("LoRA Alpha", 8, 128, 32, key="lora_alpha")
+            st.slider("LoRA Dropout", 0.0, 0.5, 0.0, key="lora_dropout")
+            st.multiselect("LoRA Targets", ["backbone", "fpn", "neck", "head"],
+                           default=["backbone", "fpn"], key="lora_targets")
+
+
+def _render_hyperparameters_tab():
+    """Training hyperparameters — only options FlashDet actually supports."""
+    st.markdown("### Training Hyperparameters")
+
+    col1, col2 = st.columns(2)
+
+    with col1:
+        with st.container(border=True):
+            st.markdown("**Core Parameters**")
+            st.slider("Epochs", 1, 500, 100, key="epochs")
+            st.select_slider("Batch Size", [2, 4, 8, 16, 32, 64, 128], value=16, key="batch_size")
+            st.select_slider("Image Size", [320, 416, 640], value=320, key="img_size",
+                             help="FlashDet default is 320; 416 and 640 also supported")
+            st.slider("Learning Rate", 1e-5, 1e-1, 1e-3, format="%.5f", key="lr")
+            st.slider("Weight Decay", 0.0, 0.1, 0.05, format="%.4f", key="weight_decay")
+
+        with st.container(border=True):
+            st.markdown("**Optimizer**")
+            st.selectbox("Optimizer", OPTIMIZERS, key="optimizer",
+                         help="AdamW (default), SGD, or MuSGD (Muon+SGD hybrid)")
+            st.checkbox("Use 8-bit Optimizer (memory saving)", value=False, key="use_8bit_optimizer")
+            st.checkbox("Gradient Clipping (max_norm=10)", value=True, key="grad_clip")
+
+    with col2:
+        with st.container(border=True):
+            st.markdown("**LR Schedule**")
+            st.caption("Cosine decay with linear warmup (fixed scheduler)")
+            st.slider("Warmup Epochs", 0, 20, 3, key="warmup_epochs")
+            st.slider("Final LR Ratio (lrf)", 0.01, 0.5, 0.1, format="%.2f", key="lr_final_ratio",
+                      help="Final LR = initial_LR × lrf")
+
+        with st.container(border=True):
+            st.markdown("**Training Options**")
+            st.checkbox("Mixed Precision (AMP FP16)", value=True, key="amp")
+            st.caption("EMA (decay=0.9998) is always enabled")
+            st.slider("DataLoader Workers", 0, 16, 4, key="num_workers")
+            st.number_input("Gradient Accumulation Steps", 1, 16, 1, key="grad_accum")
+            st.slider("Early Stopping Patience", 5, 100, 50, key="patience")
+            st.slider("Validation Interval", 1, 20, 5, key="val_interval")
+
+        with st.container(border=True):
+            st.markdown("**Multi-Scale Training**")
+            st.checkbox("Enable Multi-Scale", value=False, key="multiscale",
+                        help="Randomly varies input size between 256-416 every 10 batches")
+
+
+def _render_augmentation_tab():
+    """Data augmentation settings — only what FlashDet actually uses."""
+    st.markdown("### Data Augmentation")
+    st.caption("These are the augmentation options supported by FlashDet's training pipeline.")
+
+    col1, col2 = st.columns(2)
+
+    with col1:
+        with st.container(border=True):
+            st.markdown("**Multi-Image Augmentations** (CLI flags)")
+            st.checkbox("Mosaic (4-image composite)", value=True, key="aug_mosaic")
+            st.checkbox("MixUp (blend two images)", value=False, key="aug_mixup")
+            st.checkbox("CopyPaste (instance-level)", value=False, key="aug_copypaste")
+
+    with col2:
+        with st.container(border=True):
+            st.markdown("**Built-in Augmentations** (always applied)")
+            st.caption("These run automatically during training and are not separately configurable via CLI:")
+            st.markdown("""
+            - Random scale jitter (0.5–1.5×)
+            - Random horizontal flip
+            - Brightness / Contrast / Saturation / Hue jitter
+            - Letterbox resize (keep aspect ratio)
+            - ImageNet normalization
+            """)
+            st.info("💡 Individual geometric/color params are hardcoded in FlashDet's "
+                    "`TrainTransform`. The YAML `augment:` section exists but is not "
+                    "wired to the dataloader.")
+
+
+def _render_advanced_tab():
+    """Advanced training options — verified against FlashDet source."""
+    st.markdown("### Advanced Options")
+
+    col1, col2 = st.columns(2)
+
+    with col1:
+        with st.container(border=True):
+            st.markdown("**Loss & Assignment** (fixed in FlashDet)")
+            st.caption("These are NOT configurable — shown for reference only:")
+            st.markdown("""
+            - **Box Loss:** CIoU
+            - **Cls Loss:** BCE
+            - **Aux Loss:** L1 (on o2m head)
+            - **Assignment:** STAL (Small-Target-Aware Label)
+            - **ProgLoss:** Progressive o2m→o2o balancing
+            """)
+
+        with st.container(border=True):
+            st.markdown("**Memory & Performance**")
+            st.checkbox("Activation Checkpointing", value=False, key="activation_checkpointing",
+                        help="Trade compute for memory — useful for large models")
+            st.checkbox("Activation Offloading", value=False, key="activation_offloading")
+            st.checkbox("Optimizer in Backward", value=False, key="optimizer_in_bwd",
+                        help="Fuse optimizer step into backward pass")
+            st.checkbox("Compile Model (torch.compile)", value=False, key="compile_model")
+
+    with col2:
+        with st.container(border=True):
+            st.markdown("**Distributed Training**")
+            st.checkbox("Multi-GPU (DDP)", value=False, key="ddp")
+            if st.session_state.get("ddp"):
+                st.caption("Launch via torchrun for DDP; DataParallel as fallback")
+
+        with st.container(border=True):
+            st.markdown("**Checkpointing**")
+            st.text_input("Save Directory", value="/content/flashstudio_runs", key="save_dir")
+            st.checkbox("Save Best Only", value=True, key="save_best")
+            st.checkbox("Resume from Checkpoint", value=False, key="resume_training")
+            if st.session_state.get("resume_training"):
+                st.text_input("Checkpoint Path", placeholder="/content/runs/checkpoint_last.pth",
+                              key="resume_path")
+
+        with st.container(border=True):
+            st.markdown("**Other Training Methods**")
+            st.caption("FlashDet also supports (via separate scripts):")
+            st.markdown("""
+            - **SSL Pretraining** (BYOL / MoCo / SimCLR)
+            - **Semi-supervised** (pseudo-labeling)
+            - **Few-shot** (N-shot fine-tuning)
+            - **Active Learning** (entropy/margin query)
+            """)
+
+
+def _render_config_summary():
+    """Show configuration summary."""
+    st.markdown("### 📋 Configuration Summary")
+
+    arch_family = st.session_state.get("arch_family", "FlashDet (recommended)")
+    model = st.session_state.get("model_arch", "FlashDet-Pico")
+
+    with st.container(border=True):
+        col1, col2, col3, col4 = st.columns(4)
+        with col1:
+            st.markdown(f"**Architecture:** {arch_family}")
+            st.markdown(f"**Model:** {model}")
+            st.markdown(f"**Strategy:** {st.session_state.get('finetune_strategy', 'Full fine-tune')[:25]}")
+        with col2:
+            st.markdown(f"**Epochs:** {st.session_state.get('epochs', 100)}")
+            st.markdown(f"**Batch Size:** {st.session_state.get('batch_size', 16)}")
+            st.markdown(f"**Image Size:** {st.session_state.get('img_size', 320)}")
+        with col3:
+            st.markdown(f"**LR:** {st.session_state.get('lr', 0.001)}")
+            st.markdown(f"**Optimizer:** {st.session_state.get('optimizer', 'AdamW')}")
+            st.markdown("**Scheduler:** Cosine + warmup")
+        with col4:
+            st.markdown(f"**AMP:** {'✅' if st.session_state.get('amp') else '❌'}")
+            st.markdown("**EMA:** ✅ (always on)")
+            st.markdown(f"**Mosaic:** {'✅' if st.session_state.get('aug_mosaic') else '❌'}")
+
+    st.success("Configuration ready! Click **Next** to start training.")