| """Train on Synth6 (HOT buildings + vinyl greenhouses [+ EarthView negatives]). |
| KIND unet: from organizer baseline, L2-SP, const LR, snapshots every SNAP min. KIND siam: from Satlas init, cosine (siam_v2 recipe). |
| usage: python train13.py KIND MINUTES SNAP P_EV OUTP |
| """ |
| import math, sys, time |
| import torch, torch.nn.functional as F |
| from torch.utils.data import DataLoader |
| KIND, MINUTES, SNAP, P_EV, OUTP = sys.argv[1], float(sys.argv[2]), float(sys.argv[3]), float(sys.argv[4]), sys.argv[5] |
| sys.argv = sys.argv[:1] + ["60", "32"] |
| from synth6 import Synth6 |
| from train3 import dice |
| if KIND == "unet": |
| import segmentation_models_pytorch as smp |
| model = smp.Unet(encoder_name="resnet18", encoder_weights=None, in_channels=6, classes=3) |
| model.load_state_dict(torch.load("/workspace/unet_r18_cd.pt", map_location="cpu", weights_only=True)["state_dict"]); model = model.cuda() |
| anchor = {n: p.detach().clone() for n, p in model.named_parameters()} |
| opt = torch.optim.AdamW([{"params": list(model.parameters()), "base": 1e-4}], lr=1e-4, weight_decay=0) |
| MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).cuda(); STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).cuda() |
| bs = 64 |
| else: |
| from model_v2 import SiamCD |
| model = SiamCD("/workspace/satlas/aerial_swinb_si.pth").cuda() |
| enc = [p for n, p in model.named_parameters() if n.startswith("enc.")]; rest = [p for n, p in model.named_parameters() if not n.startswith("enc.")] |
| opt = torch.optim.AdamW([{"params": enc, "base": 1e-4}, {"params": rest, "base": 6e-4}], lr=1e-4, weight_decay=1e-4) |
| bs = 32 |
| dl = DataLoader(Synth6(10**7, P_EV), batch_size=bs, num_workers=19, pin_memory=True, persistent_workers=True, prefetch_factor=4) |
| pw = torch.tensor([2.0, 2.0]).cuda().view(1, 2, 1, 1); w3 = torch.tensor([1.0, 2.0, 2.0]).cuda() |
| t0 = time.time(); step = 0; nsnap = 1 |
| for pre, post, lab, pres, sp, sq in dl: |
| pre = pre.cuda(non_blocking=True).float() / 255; post = post.cuda(non_blocking=True).float() / 255 |
| lab, pres, sp, sq = lab.cuda(non_blocking=True), pres.cuda(non_blocking=True), sp.cuda(non_blocking=True), sq.cuda(non_blocking=True) |
| frac = min((time.time() - t0) / (MINUTES * 60), 1.0) |
| for g in opt.param_groups: |
| g["lr"] = g["base"] * (min(1, (step + 1) / 100) if KIND == "unet" else 0.5 * (1 + math.cos(math.pi * frac)) * min(1, (step + 1) / 300)) |
| if KIND == "unet": |
| y = torch.zeros_like(sp); y[lab[:, 1] > 0] = 2; y[lab[:, 0] > 0] = 1 |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| out = model(torch.cat([(pre - MEAN) / STD, (post - MEAN) / STD], 1)) |
| out = out.float(); p = out.softmax(1) |
| dl_ = sum(1 - (2 * (p[:, c] * (y == c)).sum() + 1) / (p[:, c].sum() + (y == c).sum() + 1) for c in (1, 2)) / 2 |
| l2 = sum(((q - anchor[n]) ** 2).sum() for n, q in model.named_parameters()) |
| loss = F.cross_entropy(out, y, weight=w3) + dl_ + 1e-3 * l2 |
| else: |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| ch, pr, sem_p, sem_q = model(pre, post) |
| ch, pr = ch.float(), pr.float() |
| loss = F.binary_cross_entropy_with_logits(ch, lab, pos_weight=pw) + dice(ch, lab) + 0.3 * F.binary_cross_entropy_with_logits(pr, pres) \ |
| + 0.5 * (F.cross_entropy(sem_p.float(), sp, ignore_index=255) + F.cross_entropy(sem_q.float(), sq, ignore_index=255)) |
| opt.zero_grad(set_to_none=True); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); step += 1 |
| if step % 200 == 0: print(f"step {step} {time.time() - t0:.0f}s loss {loss.item():.4f}", flush=True) |
| el = time.time() - t0 |
| if KIND == "unet" and el > nsnap * SNAP * 60: |
| torch.save({"state_dict": model.state_dict()}, f"{OUTP}_m{int(nsnap * SNAP)}.pt"); nsnap += 1; print("snapshot", int((nsnap - 1) * SNAP), flush=True) |
| if frac >= 1: break |
| torch.save({"state_dict": model.state_dict()}, OUTP + ".pt") |
| print("TRAIN_DONE", step, flush=True) |
|
|
