import streamlit as st import torch import torch.nn.functional as F from torchvision import transforms from PIL import Image import numpy as np from io import BytesIO import requests from adavit_model import AdaViTDynamic # ───────────────────────────────────────────────────────────────────────────── # PAGE CONFIG # ───────────────────────────────────────────────────────────────────────────── st.set_page_config( page_title="PatchWise — Adaptive Sparse ViT", layout="wide", initial_sidebar_state="collapsed", ) # ───────────────────────────────────────────────────────────────────────────── # GLOBAL CSS # ───────────────────────────────────────────────────────────────────────────── st.markdown(""" """, unsafe_allow_html=True) # ───────────────────────────────────────────────────────────────────────────── # DEVICE + MODEL # ───────────────────────────────────────────────────────────────────────────── device = torch.device("cpu") @st.cache_resource def load_model(): model = AdaViTDynamic( image_size=32, patch_size=4, num_classes=10, dim=256, depth=6, heads=8, mlp_dim=512, ) checkpoint = torch.load("best_model.pth", map_location=device) model.load_state_dict(checkpoint, strict=False) model.eval() return model model = load_model() CLASSES = [ "airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck" ] transform = transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), ]) # ───────────────────────────────────────────────────────────────────────────── # INFERENCE HELPER # ───────────────────────────────────────────────────────────────────────────── def run_inference(image: Image.Image): """Returns (predicted_class, confidence, keep_rate, flops_saved, patch_mask_8x8).""" tensor = transform(image).unsqueeze(0) with torch.no_grad(): outputs = model(tensor) # Unpack — model may return (logits, extras) tuple or a dict if isinstance(outputs, tuple): raw, extras = outputs[0], outputs[1] if len(outputs) > 1 else {} else: raw = outputs if isinstance(raw, dict): logits = raw["logits"] mask_flat = raw.get("mask", None) else: logits = raw mask_flat = extras.get("mask", None) if isinstance(extras, dict) else None probs = F.softmax(logits, dim=1) pred_idx = torch.argmax(probs, dim=1).item() confidence = probs[0][pred_idx].item() if mask_flat is not None: mask = mask_flat[0].cpu().numpy() keep_rate = float(mask.mean()) * 100 patch_grid = mask.reshape(8, 8) else: keep_rate = 100.0 patch_grid = np.ones((8, 8)) flops_saved = 100.0 - keep_rate return CLASSES[pred_idx], confidence * 100, keep_rate, flops_saved, patch_grid # ───────────────────────────────────────────────────────────────────────────── # SAMPLE IMAGES # ───────────────────────────────────────────────────────────────────────────── # Handpicked samples — clear, colorful, subject-centered images. # Using Wikimedia Commons with proper headers — reliably public domain. # Four demo images — multiple URL fallbacks per class so if one CDN # blocks, the next is tried. All are subject-centered, colorful images # that produce interesting patch routing maps. DEMO_SAMPLES = [ ("Frog", [ "https://images.pexels.com/photos/2062316/pexels-photo-2062316.jpeg", "https://images.pexels.com/photos/70083/frog-macro-amphibian-green-70083.jpeg", "https://images.pexels.com/photos/145683/pexels-photo-145683.jpeg", "https://images.pexels.com/photos/1490908/pexels-photo-1490908.jpeg", ]), ("Automobile", [ "https://images.pexels.com/photos/170811/pexels-photo-170811.jpeg", "https://images.pexels.com/photos/1149831/pexels-photo-1149831.jpeg", "https://images.pexels.com/photos/707046/pexels-photo-707046.jpeg", "https://images.pexels.com/photos/3802510/pexels-photo-3802510.jpeg", ]), ] @st.cache_resource(show_spinner=False) def load_demo_samples(): """For each demo class, try each fallback URL until one succeeds. Downscales to 32x32 for inference, upscales to 256x256 for display.""" import requests as _req headers = { "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 " "(KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36", "Referer": "https://www.pexels.com/", "Accept": "image/webp,image/apng,image/*,*/*;q=0.8", } samples = [] for label, urls in DEMO_SAMPLES: img_32 = None for url in urls: try: resp = _req.get(url, headers=headers, timeout=10) if resp.status_code == 200: img = Image.open(BytesIO(resp.content)).convert("RGB") img_32 = img.resize((32, 32), Image.LANCZOS) break except Exception: continue if img_32 is None: # All URLs failed — make a synthetic fallback that's still recognisable img_32 = _make_synthetic(label) img_256 = img_32.resize((256, 256), Image.NEAREST) samples.append((label, img_32, img_256)) return samples # [(label, img_32, img_256), ...] def _make_synthetic(label: str) -> Image.Image: """Generates a simple recognisable 32x32 stand-in if all URLs fail.""" from PIL import ImageDraw, ImageFilter img = Image.new("RGB", (32, 32), (20, 20, 28)) draw = ImageDraw.Draw(img) if label == "Dog": draw.rectangle([0, 0, 32, 32], fill=(200, 170, 120)) draw.ellipse([8, 8, 24, 24], fill=(160, 120, 70)) draw.ellipse([10, 5, 22, 15], fill=(170, 130, 80)) draw.ellipse([7, 4, 13, 12], fill=(130, 90, 50)) draw.ellipse([19, 4, 25, 12], fill=(130, 90, 50)) elif label == "Frog": draw.rectangle([0, 0, 32, 32], fill=(60, 120, 50)) draw.ellipse([6, 10, 26, 28], fill=(50, 160, 45)) draw.ellipse([8, 4, 24, 18], fill=(60, 175, 55)) draw.ellipse([7, 3, 14, 11], fill=(220, 220, 60)) draw.ellipse([18, 3, 25, 11], fill=(220, 220, 60)) draw.ellipse([8, 4, 13, 10], fill=(20, 70, 20)) draw.ellipse([19, 4, 24, 10], fill=(20, 70, 20)) elif label == "Ship": draw.rectangle([0, 0, 32, 18], fill=(100, 150, 210)) draw.rectangle([0, 18, 32, 32], fill=(40, 90, 160)) draw.polygon([(3, 18), (29, 18), (27, 25), (5, 25)], fill=(230, 230, 230)) draw.rectangle([8, 11, 24, 18], fill=(200, 200, 200)) draw.rectangle([12, 6, 20, 11], fill=(210, 210, 210)) draw.rectangle([15, 2, 17, 6], fill=(180, 50, 50)) elif label == "Automobile": draw.rectangle([0, 0, 32, 32], fill=(160, 185, 160)) draw.rectangle([0, 24, 32, 32], fill=(70, 70, 70)) draw.rectangle([3, 15, 29, 24], fill=(200, 40, 40)) draw.polygon([(8, 9), (24, 9), (28, 15), (4, 15)], fill=(175, 30, 30)) draw.rectangle([9, 10, 15, 14], fill=(140, 200, 230)) draw.rectangle([17, 10, 23, 14], fill=(140, 200, 230)) draw.ellipse([4, 20, 12, 28], fill=(25, 25, 25)) draw.ellipse([20, 20, 28, 28], fill=(25, 25, 25)) return img.filter(ImageFilter.GaussianBlur(0.3)) # selected_sample_idx: int index into demo_samples list, or None if "selected_sample_idx" not in st.session_state: st.session_state.selected_sample_idx = None # ───────────────────────────────────────────────────────────────────────────── # NAV # ───────────────────────────────────────────────────────────────────────────── st.markdown("""
PatchWise UCE Osmania University  ·  CSE 2025–26
""", unsafe_allow_html=True) # ───────────────────────────────────────────────────────────────────────────── # HERO # ───────────────────────────────────────────────────────────────────────────── st.markdown("""
Adaptive Sparse Vision Transformers

Not every patch
deserves attention.

PatchWise uses a reinforcement-learned policy to decide, per image, which patches are worth computing — and physically drops the rest. Same accuracy. Half the FLOPs.

""", unsafe_allow_html=True) # ── Stat strip ──────────────────────────────────────────────────────────────── st.markdown("""
54.2%
Attention FLOPs saved
79.4%
CIFAR-10 accuracy
2033fps
Jetson AGX Orin
0.49ms
Inference latency
""", unsafe_allow_html=True) # ───────────────────────────────────────────────────────────────────────────── # SECTION 01 — INFERENCE DEMO # ───────────────────────────────────────────────────────────────────────────── st.markdown("
", unsafe_allow_html=True) st.markdown("""
01 Live Sparse Inference
""", unsafe_allow_html=True) # ── Load demo samples (cached after first fetch) ───────────────────────────── with st.spinner("Fetching samples…"): demo_samples = load_demo_samples() # [(label, img_32, img_256), ...] # ── Upload row ──────────────────────────────────────────────────────────────── uploaded = st.file_uploader( "Upload", type=["png", "jpg", "jpeg"], label_visibility="collapsed", ) st.markdown("
", unsafe_allow_html=True) # ── Sample thumbnails row (full width, 4 equal columns) ────────────────────── st.markdown("""

Or try a sample — processed at CIFAR-10 resolution (32×32)

""", unsafe_allow_html=True) s_cols = st.columns(len(demo_samples), gap="medium") for i, (label, img_32, img_256) in enumerate(demo_samples): with s_cols[i]: st.image(img_256, use_container_width=True) if st.button(label, key=f"sample_{i}"): st.session_state.selected_sample_idx = i uploaded = None # ── Resolve active image ────────────────────────────────────────────────────── active_image = None active_label = None if uploaded is not None: active_image = Image.open(uploaded).convert("RGB") active_label = "uploaded" st.session_state.selected_sample_idx = None elif st.session_state.selected_sample_idx is not None: idx = st.session_state.selected_sample_idx active_label = demo_samples[idx][0] active_image = demo_samples[idx][1] # use the 32×32 version for inference # ── Inference output ────────────────────────────────────────────────────────── st.markdown("
", unsafe_allow_html=True) if active_image is not None: pred_class, confidence, keep_rate, flops_saved, patch_grid = run_inference(active_image) out_img_col, out_patch_col, out_result_col = st.columns([1, 1, 1.4], gap="large") with out_img_col: st.markdown('

Input image

', unsafe_allow_html=True) st.image(active_image, use_container_width=True) with out_patch_col: st.markdown('

Active patch map

', unsafe_allow_html=True) # Render patch grid as a colour-coded image patch_vis = np.zeros((patch_grid.shape[0], patch_grid.shape[1], 3), dtype=np.uint8) patch_vis[patch_grid == 1] = [56, 189, 248] # cyan-ish = kept patch_vis[patch_grid == 0] = [18, 18, 22] # near-black = pruned patch_img = Image.fromarray(patch_vis).resize((128, 128), Image.NEAREST) st.image(patch_img, use_container_width=True) with out_result_col: st.markdown(f"""
Prediction
{pred_class}
Confidence
{confidence:.1f}%
Patches kept {keep_rate:.1f}%
FLOPs saved {flops_saved:.1f}%
Routing policy A3C · EMA-controlled
""", unsafe_allow_html=True) else: st.markdown("""

Upload an image or select a sample above
to observe adaptive patch routing live.

""", unsafe_allow_html=True) # ───────────────────────────────────────────────────────────────────────────── # SECTION 02 — METHODOLOGY # ───────────────────────────────────────────────────────────────────────────── st.markdown("
", unsafe_allow_html=True) st.markdown('
', unsafe_allow_html=True) st.markdown("
", unsafe_allow_html=True) st.markdown("""
02 Three-Phase Development
Phase 01
Gumbel-Softmax differentiable approximation
Embedded a lightweight decision network into DeiT-Small. Used the Gumbel-Softmax reparameterization trick to allow gradient flow through discrete keep/drop decisions. Dual-objective loss: cross-entropy + MSE usage penalty toward a target keep rate.
Training–inference gap · soft probs vs hard actions
Phase 02
Hybrid single-agent A2C with Bernoulli policy
Transitioned to true discrete decisions. Actor outputs one logit per patch sampled via a Bernoulli distribution. Patches physically removed before attention blocks using x.detach() gating. Critic reduces variance via advantage estimation.
High variance · correlated sequential data · gradient conflict
Phase 03
Asynchronous A3C · Categorical policy · EMA controller
Four parallel Hogwild! workers decorrelate training data. Categorical distribution with two independent logits per patch eliminates gradient tug-of-war. EMA controller enforces keep-rate budget. Attention-guided dense rewards from early transformer layers. Curriculum warmup stabilises joint backbone + policy training.
Stable convergence · 79.4% accuracy @ ≈47.7% keep rate
""", unsafe_allow_html=True) # ───────────────────────────────────────────────────────────────────────────── # SECTION 03 — CONTRIBUTIONS + ARCHITECTURE (side by side) # ───────────────────────────────────────────────────────────────────────────── st.markdown("
", unsafe_allow_html=True) st.markdown('
', unsafe_allow_html=True) st.markdown("
", unsafe_allow_html=True) contrib_col, arch_col = st.columns([1, 1], gap="large") with contrib_col: st.markdown("""
03 Contributions
01 Attention-Guided Reward Mechanism — dense, spatially-aware feedback extracted from early transformer layers, replacing sparse end-of-episode signals.
02 EMA Controller — exponential moving average of batch-wise keep rate enforces a strict computational budget, eliminating keep-rate collapse.
03 Curriculum Warmup — RL policy disabled during early epochs so the backbone learns stable representations before pruning activates.
04 Categorical → Bernoulli upgrade — dual-logit formulation decouples accuracy and sparsity gradients, resolving the Phase 2 tug-of-war.
05 Hogwild! A3C parallelisation — four asynchronous workers update a shared global model, decorrelating data and accelerating convergence.
06 Edge validation on NVIDIA Jetson AGX Orin — 2033 FPS, 0.49 ms latency, confirming real-world deployment viability.
""", unsafe_allow_html=True) with arch_col: st.markdown("""
04 Model Configuration
BackboneCustom Vision Transformer (DeiT-inspired)
RL ControllerAsynchronous A3C · 4 workers
Policy distributionCategorical (Softmax, 2 logits / patch)
DatasetCIFAR-10 · CIFAR-100 stress test
Input resolution32 × 32 px
Patch size4 × 4 → 64 tokens total
Embedding dim256
Depth6 transformer layers
Attention heads8
Target keep rate50% (EMA-enforced)
Edge hardwareNVIDIA Jetson AGX Orin
CIFAR-100 stability~57% accuracy @ 300 epochs
""", unsafe_allow_html=True) # ───────────────────────────────────────────────────────────────────────────── # FOOTER # ───────────────────────────────────────────────────────────────────────────── st.markdown("""
PatchWise  ·  Adaptive Sparse Vision Transformers for Edge AI Srivarsha S · Aila Kaushik · Adla Sahithi  —  UCE Osmania University · CSE · May 2026
""", unsafe_allow_html=True)