| """T-SRDA (PDAViT dual-attention, temporal axis) configuration. | |
| Data protocol is now byte-for-byte the OFFICIAL STCLN PASTIS protocol | |
| (XiaoleiQinn/STCLN, mirrored at ../test/stcln_ref/PASTIS), so tsrda is | |
| directly comparable to the det/DaViT and PhenoProto-SSL arms. | |
| [OFFICIAL] = verbatim from the reference implementation. Changing these | |
| breaks comparability with the published 0.4843 mIoU. | |
| [TSRDA] = this arm's architectural variable, or a deliberate divergence | |
| (each one is called out in the comment). | |
| Fold budget — all four roles are fold-disjoint: | |
| pretrain fold 5 496 patches unlabelled, masked reconstruction | |
| train fold 1 76 IDs labelled, 2 crops each -> 152 | |
| val fold 2 76 IDs labelled, 2 crops each -> 152 | |
| test fold 4 482 patches full 128x128, no cropping | |
| fold 3 474 patches unused by the protocol | |
| """ | |
| import os | |
| from pathlib import Path | |
| # ─── Paths ────────────────────────────────────────────────────────────────── | |
| # Derived from this file's own location so the project moves between machines | |
| # without edits. Current layout on this box: | |
| # | |
| # /home/ubuntu/tsrda/ <- REPO_ROOT = EXP_ROOT (code + checkpoints) | |
| # /home/ubuntu/test/PASTIS/ <- PASTIS_ROOT (dataset) | |
| # | |
| # PASTIS_ROOT is searched in the order below; the first existing hit wins. | |
| # Override either root by env var, e.g. | |
| # PASTIS_ROOT=/mnt/data/PASTIS python3 check_data.py | |
| REPO_ROOT = Path(__file__).resolve().parent | |
| _PASTIS_CANDIDATES = ( | |
| REPO_ROOT / "PASTIS", # dataset inside the repo | |
| REPO_ROOT.parent / "PASTIS", # sibling of the repo | |
| REPO_ROOT.parent / "test" / "PASTIS", # this box: /home/ubuntu/test/PASTIS | |
| ) | |
| def _find_pastis_root(): | |
| env = os.environ.get("PASTIS_ROOT") | |
| if env: | |
| return Path(env) | |
| for cand in _PASTIS_CANDIDATES: | |
| if (cand / "metadata.geojson").is_file(): | |
| return cand | |
| raise FileNotFoundError( | |
| "PASTIS dataset not found. Looked for metadata.geojson in:\n " | |
| + "\n ".join(str(c) for c in _PASTIS_CANDIDATES) | |
| + "\nSet the PASTIS_ROOT environment variable to the dataset directory." | |
| ) | |
| PASTIS_ROOT = _find_pastis_root() | |
| DATA_S2_DIR = PASTIS_ROOT / "DATA_S2" | |
| ANNOT_DIR = PASTIS_ROOT / "ANNOTATIONS" | |
| META_PATH = PASTIS_ROOT / "metadata.geojson" | |
| NORM_PATH = PASTIS_ROOT / "NORM_S2_patch.json" | |
| EXP_ROOT = Path(os.environ.get("EXP_ROOT", REPO_ROOT)) | |
| CKPT_DIR_PRE = EXP_ROOT / "checkpoints" / "pretrain" | |
| CKPT_DIR_FT = EXP_ROOT / "checkpoints" / "finetune" | |
| # ─── Reproducibility ──────────────────────────────────────────────────────── | |
| SEED = 3407 # [OFFICIAL] torch.manual_seed in | |
| # finetuning_STCLN.py:25 and test_STCLN.py:11 | |
| VARIANCE_SEEDS = (3407, 42, 1234) | |
| # [OFFICIAL] the three seeds used for | |
| # seed-variance runs. Pass with --seed. | |
| # ─── Model dimensions ─────────────────────────────────────────────────────── | |
| N_CHANNELS = 10 # [OFFICIAL] | |
| D_MODEL = 256 # [OFFICIAL] | |
| N_HEADS = 8 # [OFFICIAL] | |
| N_CLASSES = 20 # [OFFICIAL] 0=background, 19=void, 1-18 scored | |
| T_PAD = 48 # inert; UTAE_ASPP absorbs it for constructor | |
| # compatibility and the encoder self-pads | |
| # T-SRDA schedule: KV reduction ratio per block (replaces Swin/DaViT windows). | |
| R_SCHEDULE = [4, 4, 4] # [TSRDA] the architectural variable | |
| # Legacy temporal-Swin params — kept ONLY so the train scripts' constructor | |
| # call sites compile unchanged. TemporalSRDAEncoder ignores them. | |
| TSE_WINDOWS = [4, 8, 16] | |
| TSE_SHIFTS = [0, 4, 0] | |
| # ─── Protocol: geometry ───────────────────────────────────────────────────── | |
| PATCH_SIZE = 128 # [OFFICIAL] PASTIS native patch | |
| CROP_SIZE = 32 # [OFFICIAL] PATCH_SIZE // 4 | |
| REF_DATE_STR = "2018-09-01" # [OFFICIAL] | |
| USE_INDEX_POSITIONS = True # [OFFICIAL] the model is fed torch.arange(T), | |
| # NOT real day offsets. See | |
| # finetuning_STCLN.py:198 — every forward call | |
| # passes `torch.tensor(range(T))`. Real days are | |
| # still returned by the dataset for analysis. | |
| # ─── Protocol: pretrain split ─────────────────────────────────────────────── | |
| PRETRAIN_FOLDS = [5] # [OFFICIAL] 496 patches | |
| PRETRAIN_CROP_GRID = 4 # [OFFICIAL] inner `for i in range(4): for j in | |
| # range(4)` loop -> 16 crops per patch, each a | |
| # SEPARATE optimizer step. | |
| # 496 / 4 = 124 batches x 16 = 1,984 steps/ep | |
| CROPS_PER_PATCH = PRETRAIN_CROP_GRID ** 2 # 16 | |
| # ─── Protocol: finetune splits (hardcoded patch IDs, not a fold filter) ───── | |
| # The official script builds the dataset with folds=[1] / folds=[2] and then | |
| # OVERWRITES dataset.id_patches with these literal lists. Verified against | |
| # metadata.geojson: every train ID is fold 1, every val ID is fold 2, the two | |
| # sets are disjoint, and neither touches the fold-4 test set or fold-5 | |
| # pretrain pool. Duplicates are intentional and are KEPT — they make those | |
| # patches count twice per epoch. | |
| FT_CROP_IJ = [(0, 0), (1, 1)] # [OFFICIAL] finetuning_STCLN.py:185 — only the | |
| # first two diagonal crops of the 4x4 grid. | |
| # 76 / 2 = 38 batches x 2 crops = 76 steps/ep | |
| TRAIN_PATCH_IDS = [ | |
| 10279, 40282, 10174, 40538, 10457, 30088, 10151, 30602, 40275, 20052, | |
| 20153, 10056, 20147, 40420, 10450, 20614, 20167, 40446, 20251, 20103, | |
| 20613, 20349, 20384, 40337, 40163, 10030, 10068, 10399, 30327, 10289, | |
| 30351, 40198, 20167, 20153, 20207, 20203, 10000, 10354, 20214, 10007, | |
| 10392, 30055, 10054, 20009, 10040, 10147, 10110, 10105, 20398, 20469, | |
| 20203, 20254, 30402, 30013, 40200, 30689, 20115, 20356, 20345, 20466, | |
| 40033, 10153, 30153, 30012, 40340, 30140, 40214, 10129, 30211, 20078, | |
| 30109, 30013, 20116, 20243, 20085, 40271, | |
| ] # [OFFICIAL] 76 entries, 72 unique (20153, 20167, 20203, 30013 appear 2x) | |
| VAL_PATCH_IDS = [ | |
| 10127, 40438, 30643, 30619, 30436, 10306, 10127, 30226, 40027, 40279, | |
| 40085, 40455, 40426, 40444, 40298, 40231, 20498, 40298, 20149, 20315, | |
| 20340, 20618, 20412, 20231, 40006, 40138, 40186, 10087, 40430, 30194, | |
| 10247, 40383, 20364, 20128, 20060, 20118, 40001, 20080, 40063, 30000, | |
| 40063, 40000, 40348, 10021, 40001, 20023, 30000, 20441, 20235, 20344, | |
| 20297, 20576, 30014, 30275, 30298, 30083, 20206, 20443, 20423, 20380, | |
| 30064, 30084, 30145, 10107, 10225, 40558, 10383, 20131, 30010, 20022, | |
| 30295, 20042, 20031, 40046, 20358, 30142, | |
| ] # [OFFICIAL] 76 entries, 71 unique (10127, 40298, 40001, 40063, 30000 2x) | |
| # ─── Protocol: test split ─────────────────────────────────────────────────── | |
| TEST_FOLDS = [4] # [OFFICIAL] 482 patches, spatially disjoint | |
| # from folds 1/2/5 with a 1 km buffer | |
| TEST_FULL_PATCH = True # [OFFICIAL] test_STCLN.py:55 does | |
| # `split = W // 1` -> the whole 128x128 patch | |
| # is evaluated, NOT 32x32 crops. | |
| # 482 patches ~ 482 x 16 = 7,712 crop-equiv | |
| # ─── Pretrain hyperparameters ─────────────────────────────────────────────── | |
| PRE_EPOCHS = 100 # [OFFICIAL] | |
| PRE_BATCH = 4 # [OFFICIAL] PATCHES per batch, not crops. The | |
| # inner 4x4 loop then does 16 separate | |
| # optimizer steps on 4 crops each. Raising this | |
| # changes the optimization trajectory, not just | |
| # memory — do not tune it for the headline run. | |
| PRE_LR = 1e-4 # [OFFICIAL] flat, no schedule | |
| PRE_WD = 0.0 # [OFFICIAL] AdamW with wd=0 reduces to Adam | |
| # exactly. Set to 0.0 (from 0.01) so the only | |
| # variable vs A_linear is the temporal encoder. | |
| PRE_CLIP = 5.0 # [OFFICIAL] | |
| MASK_RATIO = 0.4 # [OFFICIAL] | |
| NDVI_THRESH = 0.2 # [OFFICIAL] | |
| CLOUD_GATE = 0.9 # [OFFICIAL] STCLN.py:200 `mask[clusterLmean | |
| # <= 0.9] = 1`. clusterLmean averages the NDVI | |
| # indicator over H and W, so this is a PER-FRAME | |
| # gate: a timestep whose frame is <=90% | |
| # vegetated is left ENTIRELY visible. It is NOT | |
| # a per-pixel vegetation filter. Measured on 20 | |
| # fold-5 patches: fires on 90.6% of frames, | |
| # leaving 96.2% of all values visible. | |
| PRE_SAVE_EVERY = 20 # [TSRDA] keep a numbered milestone every N | |
| # epochs (0, 19, 39, 59, 79, 99 at N=20). | |
| # Independent of this, `latest.tar` is written | |
| # EVERY epoch with optimizer + scaler + RNG | |
| # state, so a crash costs at most one epoch | |
| # (~26 min) rather than up to N (~8.7 h). | |
| DEEP_SUP = False # [OFFICIAL] the deep-supervision term is | |
| # commented out upstream | |
| # (pretraining_STCLN.py:141 | |
| # `l = loss(output, target)#+loss(x2,target)`). | |
| # ─── Finetune hyperparameters ─────────────────────────────────────────────── | |
| FT_EPOCHS = 100 # [OFFICIAL] | |
| FT_BATCH = 2 # [OFFICIAL] PATCHES per batch (x2 crop-steps) | |
| FT_LR = 1e-4 # [OFFICIAL] flat Adam, NO scheduler | |
| FT_WD = 0.0 # [OFFICIAL] AdamW with wd=0 == Adam | |
| FT_AUGMENT = False # [OFFICIAL] no augmentation in finetuning_STCLN | |
| FT_SCHEDULER = None # [OFFICIAL] flat LR, no schedule. | |
| FT_PATIENCE = 0 # [OFFICIAL] no early stopping (grep: 0 matches | |
| # in the reference). 0 disables it entirely. | |
| # MUST stay 0: the primary report is the | |
| # FIXED-EPOCH ep99 number. Early stopping on | |
| # 152 val crops manufactured a 147x variance | |
| # effect that vanished at fixed epoch, and if | |
| # it fires at ep58 there is no ep99 checkpoint | |
| # to compare against A_linear at all. | |
| # ─── Evaluation ───────────────────────────────────────────────────────────── | |
| EVAL_BATCH = int(os.environ.get("EVAL_BATCH", 4)) | |
| # [HW] official test_STCLN.py uses 1. Full-patch | |
| # inference is numerically identical at any | |
| # batch size (nothing in eval mode depends on | |
| # the batch), so this is a pure memory/speed | |
| # knob with no effect on the reported metrics. | |
| # MEASURED for T-SRDA on this 23.6 GB L4 against | |
| # the worst case (fold 4 runs to T=61): | |
| # batch 1 -> 5.05 GB peak 482 batches | |
| # batch 2 -> 9.50 GB peak 241 batches | |
| # batch 4 -> 16.86 GB peak 121 batches | |
| # Default 4 gives the protocol's 121 batches and | |
| # fits with ~6.7 GB spare ON A FREE CARD. | |
| # asr-gpu.service is systemd-`enabled` on this | |
| # box and takes ~8 GB when it comes back after a | |
| # reboot, leaving ~15.6 GB — under batch 4. If | |
| # eval OOMs, just rerun with EVAL_BATCH=2; the | |
| # numbers are unchanged. | |
| EVAL_TTA = False # [OFFICIAL] test_STCLN.py has no TTA | |
| # ─── PASTIS class nomenclature ────────────────────────────────────────────── | |
| # THE single source of truth for this repo. The list that shipped in the | |
| # original tsrda scripts was wrong for 12 of 20 classes (it invented | |
| # "Protein peas", "Flax", "Sugar beet", "Vineyards", "Nuts" and mislabelled | |
| # 13 as Soy / 17 as Potatoes). Verified identical to PhenoProto-SSL's | |
| # splits.PASTIS_CLASSES, which preflight.py asserts against the official | |
| # PASTIS benchmark list. Never redefine this locally in another module. | |
| PASTIS_CLASSES = { | |
| 0: "Background", | |
| 1: "Meadow", | |
| 2: "Soft winter wheat", | |
| 3: "Corn", | |
| 4: "Winter barley", | |
| 5: "Winter rapeseed", | |
| 6: "Spring barley", | |
| 7: "Sunflower", | |
| 8: "Grapevine", | |
| 9: "Beet", | |
| 10: "Winter triticale", | |
| 11: "Winter durum wheat", | |
| 12: "Fruits, vegetables, flowers", | |
| 13: "Potatoes", | |
| 14: "Leguminous fodder", | |
| 15: "Soybeans", | |
| 16: "Orchard", | |
| 17: "Mixed cereal", | |
| 18: "Sorghum", | |
| 19: "Void label", | |
| } | |