File size: 5,456 Bytes
1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 fc5c00d 1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """
Training protocol shared by every baseline / ablation run, and the variants.
All variants are trained FROM SCRATCH with the same two-phase protocol, split, number of epochs and
checkpoint-selection rule, so that differences are attributable to the ablated component only.
The two phases mirror the staged training of the released model:
phase 1 localisation: weighted MSE + soft Dice (lambda_rec 100), lr 1e-4 (as train_pix2pix.ipynb)
phase 2 magnitude: per-image WFM-RME (lambda_rec 10) + force consistency (lambda_force 100),
lr 1e-5, initialised from the best phase-1 generator (as train_pix2pix_settings_mean.ipynb)
WFM-RME alone does not learn to localise forces from random weights (tested: pixel r ~ 0 after 100 epochs),
which is why the released model, too, was first trained with the weighted MSE + Dice loss.
Unlike the released model, every variant sees its full input (incl. substrate channels) from the start.
"""
import copy
PROTOCOL = dict(
train_manifest="alltune_train", # 280 pairs, released model's training split
val_manifest="alltune_val", # 81 pairs, used for checkpoint selection only
phases=[
dict(name="phase1", epochs=100, rec_loss="wmd", lambda_rec=100.0, lambda_force=0.0, g_lr=1e-4, d_lr=1e-4),
dict(name="phase2", epochs=40, rec_loss="rme", lambda_rec=10.0, lambda_force="from_variant", g_lr=1e-5, d_lr=1e-5),
],
batch_size=2,
img_size=1024,
augment=True, # data/augmentations.AdvancedAugmentations
beta1=0.5, beta2=0.999,
warmup_frac=0.1, # cosine schedule with linear warm-up (per phase)
lambda_gan=1.0, # vanilla (BCE) PatchGAN loss, both phases
lambda_force=100.0, # phase 2: MSE between mean predicted and mean target intensity
arch="s2f", in_channels=3, use_multi_scale_input=True, bridge="cbam", attention_gates=True,
substrate_mode="legacy", # see experiments/manifest_data.py
input_mode="bf",
seed=0,
)
VARIANTS = {
# reference
"s2f_full": dict(),
# baselines
"unet": dict(arch="unet", lambda_gan=0.0), # plain U-Net, no attention, no discriminator
"unet_gan": dict(arch="unet"), # pix2pix-style U-Net + PatchGAN
"mask_only": dict(input_mode="mask"), # binary cell mask instead of bright field
# loss ablations
"s2f_no_gan": dict(lambda_gan=0.0),
"s2f_no_force": dict(lambda_force=0.0, share_phase1_with="s2f_full"), # force term is a phase-2 term only
"s2f_gan_0.1": dict(lambda_gan=0.1),
"s2f_gan_5": dict(lambda_gan=5.0),
# input ablations
"s2f_no_substrate": dict(in_channels=1, substrate_mode="none"),
"s2f_nominal_substrate": dict(substrate_mode="nominal"),
# architecture ablations
"s2f_no_multiscale": dict(use_multi_scale_input=False),
"s2f_bridge_no_cbam": dict(bridge="cnn"),
"s2f_bridge_plain": dict(bridge="plain"),
"s2f_no_attention_gates": dict(attention_gates=False),
# image-only (no substrate channels) counterparts: the substrate channels did not help (2 seeds), so the
# architecture and adversarial-loss comparisons are repeated on the image-only reference 's2f_no_substrate'
"unet_no_substrate": dict(arch="unet", lambda_gan=0.0, in_channels=1, substrate_mode="none"),
"unet_gan_no_substrate": dict(arch="unet", in_channels=1, substrate_mode="none"),
"s2f_no_substrate_no_gan": dict(in_channels=1, substrate_mode="none", lambda_gan=0.0),
"s2f_no_substrate_no_multiscale": dict(in_channels=1, substrate_mode="none", use_multi_scale_input=False),
"s2f_no_substrate_bridge_plain": dict(in_channels=1, substrate_mode="none", bridge="plain"),
"s2f_no_substrate_no_attention_gates": dict(in_channels=1, substrate_mode="none", attention_gates=False),
}
# Recommended order if GPU time is limited (S2F about 7 h per run on an RTX A6000, U-Net about 2.5 h)
PRIORITY = ["s2f_full", "unet", "unet_gan", "s2f_no_gan", "s2f_no_substrate", "s2f_no_force",
"s2f_no_attention_gates", "s2f_no_multiscale", "s2f_bridge_plain",
"s2f_bridge_no_cbam", "s2f_nominal_substrate", "s2f_gan_0.1", "s2f_gan_5"]
# 'mask_only' needs a cell mask (cell_mask.tif) in every sample folder and is therefore not in PRIORITY.
def get_config(variant, epochs=None, g_lr=None, d_lr=None, **overrides):
"""epochs / g_lr / d_lr override every phase (smoke tests); other overrides are top-level keys."""
if variant not in VARIANTS:
raise KeyError(f"unknown variant '{variant}'. Options: {', '.join(VARIANTS)}")
cfg = copy.deepcopy(PROTOCOL)
cfg.update(VARIANTS[variant])
cfg.update({k: v for k, v in overrides.items() if v is not None})
cfg["variant"] = variant
if cfg["substrate_mode"] == "none":
cfg["in_channels"] = 1
for ph in cfg["phases"]:
if ph["lambda_force"] == "from_variant":
ph["lambda_force"] = cfg["lambda_force"]
for k, v in (("epochs", epochs), ("g_lr", g_lr), ("d_lr", d_lr)):
if v is not None:
ph[k] = v
cfg["epochs"] = sum(ph["epochs"] for ph in cfg["phases"]) # total, used to detect finished runs
return cfg
|