FP4S Floor Plan Segmentation: CCT Ablation Study

# FP4S reproduction — smoke scope + CCT ablation Reproducing **"Floor Plan Image Segmentation Via Scribble-Based Semi-Weakly Supervised Learning"** (Chen et al., SSRN 4727643) with the official code (`github.com/JanineCHEN/FP4S`) and the original Harvard Dataverse dataset (`doi:10.7910/DVN/MDIRHE`). ## Scope (agreed) | | | |---|---| | Unlabeled pool | **1,000** ArchDaily images (paper's 1k ablation point) | | Epochs | **20** (paper: 200) | | Iters/epoch | **200** (paper: `len(unsupervised_loader)` = 500) | | Config | repo defaults — resnet18, bs=2, 512 px, AdamW lr=1e-4 wd=0.1, BCE, no cutmix, no pretrain, no input norm, `w_sm=w_edge=1.0` | | Arms | **`fp4s-with-cct`** (full method) vs **`fp4s-without-cct`** (paper Table 5 ablation), forked as parallel branches |

## Stage 1 — Environment, scope and two of the five fixes Sets up the run and pins everything that must be reproducible. - **Config comes from the repo itself.** `tools/config.py:get_config()` is an `argparse` parser that the repo calls *at module import time* (in `models.py`, `utils.py`, `loss.py`). So `sys.argv` is set **before** the first repo import — otherwise the notebook's own arguments would be parsed as FP4S flags. - **Scope**: 1,000 unlabeled images, 20 epochs × 200 iterations = 4,000 iterations per arm. - **Seeds** (`random`, `numpy`, `torch`, CUDA, `PYTHONHASHSEED`) are set before any model is built, so both arms start from the same weights and see the same labeled batches. - **Fix 2 — `set_detect_anomaly(False)`.** `main.py` ships anomaly detection *enabled*. Measured cost: 1.208 s/iter with it on vs 0.387 s/iter with it off — a **3.1× slowdown** for a debugging feature. - **Fix 5 — TF32 on.** Free tensor-core path for fp32 matmul/conv on this A100. The cell also asserts a GPU is present (never silently train on CPU) and prints the file count of every dataset directory, so a missing folder fails here rather than 20 minutes into training.

# === Cell 1: environment, config, fixes 2 & 5 =================================
import os, sys, json, time, random, io, contextlib
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F

REPO, RUNS = '/home/user/repos/FP4S', '/home/user/fp4s_runs'
os.makedirs(RUNS, exist_ok=True)
os.chdir(REPO)
if REPO not in sys.path:
    sys.path.insert(0, REPO)

# ---- Reproduction scope (smoke) ---------------------------------------------
NUM_UNSUP, NUM_EPOCHS, ITERS_PER_EPOCH = 1000, 20, 200
BATCH, TRAINSIZE, SEED = 2, 512, 0

# The repo builds its config with argparse AT IMPORT TIME (tools/config.py:get_config()
# is called at module scope in models.py / utils.py / loss.py), so sys.argv must be set first.
sys.argv = ['fp4s', '--seed', str(SEED), '--num_epochs', str(NUM_EPOCHS),
            '--num_unsupervised_imgs', str(NUM_UNSUP), '--batchsize', str(BATCH),
…

torch 2.7.1+cu128 | NVIDIA A100-SXM4-40GB | 42 GB
scope   : unsup=1000 epochs=20 iters/epoch=200 -> 4000 iters/arm
cfg     : backbone=resnet18 bs=2 size=512 lr=0.0001 classes=25 channel=32
cfg abl : cutmix=False abCE=False focal=False norm=False pretrain=False w_sm=1.0 w_edge=1.0
  ./dataset/FP_train/img            558 files
  ./dataset/FP_train/gt             558 files
  ./dataset/FP_train/mask           558 files
  ./dataset/FP_train/gray           558 files
  ./dataset/FP_train/edge…

## Stage 2 — Download the unlabeled half of the dataset The Dataverse archive ships the labeled data, but the unlabeled pool is only a **list of 68,316 ArchDaily URLs** (`unsupervised_list.csv`) — the images have to be fetched. This mirrors what `main.py` does (seeded shuffle → `unsupervised_{i}.jpg` → an `unAnnotated_FPs_1000.txt` manifest) with two practical changes: - **16 parallel workers** instead of a serial loop (measured ~0.05 s/image serially). - **Verify and replace.** Each file is opened with PIL to confirm it actually decodes; dead URLs are replaced from further down the shuffled list, so we end up with exactly 1,000 *usable* images rather than 1,000 attempts. Result: 1,000 images in **54 s**. Downloading is not a bottleneck at this scale — but it would be ~9 minutes for the paper's best 10k setting.

# === Cell 2: fetch the 1,000 unlabeled ArchDaily images =======================
# Mirrors main.py's download block (seeded shuffle of unsupervised_list.csv -> unsupervised_{i}.jpg
# + an unAnnotated_FPs_{N}.txt manifest), but threaded (measured ~0.05 s/img serial) and
# with decode verification + replacement of dead URLs.
import pandas as pd, requests
from concurrent.futures import ThreadPoolExecutor
from PIL import Image

UNSUP_DIR = './dataset/FP_unsupervised/'
LIST_PATH = cfg.image_unsupervised_list_path      # ./dataset/unAnnotated_FPs_1000.txt
CSV_PATH  = '/home/user/dv/unsupervised_list.csv' # from the Dataverse archive
os.makedirs(UNSUP_DIR, exist_ok=True)

urls = list(pd.read_csv(CSV_PATH)['img_url'].astype(str))
random.Random(SEED).shuffle(urls)
print(f'{len(urls)} candidate URLs; need {NUM_UNSUP}')

sess = requests.Session()
…

68316 candidate URLs; need 1000

## Stage 3 — Add the ablation switch, and make evaluation fast Two jobs here. **1. The `use_cct` switch.** The baseline arm (paper Table 5, "w/o CCT") is FP4S with the cross-consistency mechanism removed: no unlabeled forward pass, no 30 auxiliary decoders, no MSE consistency term — just the supervised scribble + edge + smoothness losses. Rather than copy-paste a second model class, we insert an early `return` into `nets/models.py:FP4S.forward` right before the unsupervised block. The patch is applied **before** the first import of `nets.models`, so no module-reload tricks are needed, and it is idempotent (re-running the cell detects its own marker). **2. Fix 4 — a vectorized `evaluation()`.** The repo's metric loops over 25 classes in Python on the CPU, costing **24 s per validation pass** — which across 40 validation passes would have been a large share of this run. The replacement does the same arithmetic as batched GPU tensor ops: per-class min-max normalize the raw logits, threshold at 0.9, then Dice/IoU from TP/FP/FN, including the repo's "both empty → 1.0" convention. The cell then **asserts equivalence** against the original on random tensors. Result: max absolute difference **0.0** for both Dice and IoU — identical numbers, ~147× faster. This is why the speedup is safe to claim: it is a measured identity, not an approximation.

# === Cell 3: patch the ablation switch into FP4S, then import the repo =========
# The "w/o CCT" arm (paper Table 5) = FP4S with the cross-consistency mechanism removed:
# no unlabeled forward pass, no 30 aux decoders, no MSE consistency term — supervised
# scribble + edge + smoothness losses only. We add a `use_cct` flag by inserting an early
# return into nets/models.py:FP4S.forward, right before the unsupervised block. Patching
# happens BEFORE importing nets.models so no reload games are needed.
MODELS_PY = os.path.join(REPO, 'nets/models.py')
MARKER    = '# --- FP4S w/o CCT ablation (paper Table 5)'
ANCHOR    = '        # Get unsupervised predictions\n'
PATCH = f'''        {MARKER} ---
        if not getattr(self, 'use_cct', True):
            _sup = sal_loss1 + edge_loss + sal_loss2
            _cl = {{'sal_loss1': sal_loss1.item(), 'edge_loss': edge_loss.item(),
                   'sal_loss2': sal_loss2.item(), 'smoothLoss_cur1': float(smoothLoss_cur1),
                   'smoothLoss_cur2': float(smoothLoss_cur2), 'unsup_w': 0.0,
                   'unsup_loss': 0.0, 'sup_loss': _sup.item(),
                   'total_loss': _sup.item(), 'total_loss_mean': _sup.mean().item()}}
            return _sup.mean(), _cl, {{'sal1': sal_init, 'edge': edge_map, 'sal2': sal_ref}}
…

patched nets/models.py with the use_cct switch

## Stage 4 — Data loaders and the training driver One function, `train_arm(tag, use_cct)`, runs either arm. Both arms share it, so any difference in the results comes from the mechanism and not from two diverging training scripts. **Augmentation.** The same `transf_aug` pipeline as `main.py` (flips, ±10° rotation, random resized crop). It is applied to image / gt / mask / gray / edge with the **RNG state reset before each** so all five get exactly the same geometric transform — otherwise the labels would no longer align with the image. **Fix 1 — `clamp(0, 1)`, the blocker.** Bilinear resize and rotation of a *binary* map undershoot slightly: the `edge` tensor comes out with a minimum of about `-3.6e-07`. `nn.BCELoss` validates that targets lie in [0, 1] and fires a device-side assert on **iteration 0**. Without this clamp the released code cannot train at all on modern PyTorch. This was the single hardest problem to diagnose, because the CUDA assert surfaces as an unrelated-looking error far from its cause. **Fix 3 — `num_workers=8`.** The repo's loaders use `num_workers=1`, which delivers a batch every 1.84 s against a 0.39 s GPU step — the GPU sits idle ~80% of the time. With 8 workers it is 0.193 s/batch and training becomes GPU-bound. Measured: 1 → 1.841 s, 4 → 0.510 s, 8 → **0.193 s**, 12 → 0.271 s (over-subscribed). **What each epoch does**: 200 iterations over the labeled set (cycled, since 558 images ≫ 200×2 samples), then a full validation pass over all 28 val i…

# === Cell 4: loaders + training/validation driver (shared by both branches) ====
from itertools import cycle
import torchvision.transforms as transforms

# Identical to main.py's transf_aug; applied on-GPU with a shared RNG state so that
# image/gt/mask/gray/edge get exactly the same geometric transform.
transf_aug = transforms.Compose([
    transforms.RandomHorizontalFlip(0.5),
    transforms.RandomVerticalFlip(0.5),
    transforms.RandomRotation(10),
    transforms.RandomResizedCrop((cfg.trainsize, cfg.trainsize), scale=(0.7, 1.0)),
])

def build_loaders(need_unsup, num_workers=8):
    """FIX 3: repo uses num_workers=1 (1.84 s/batch, data-bound); 8 -> 0.193 s/batch."""
    sup = tdata.DataLoader(
        SribbleObjDataset('./dataset/FP_train/img/', './dataset/FP_train/gt/', './dataset/FP_train/mask/',
                          './dataset/FP_train/gray/', './dataset/FP_train/edge/', TRAINSIZE, cfg.ifNorm),
…