| """Fine-tune on Synth5 (real JL1 + synthetic). KIND siam: from siam_v2; KIND unet: from soup_all with L2-SP. |
| usage: python train12.py KIND MINUTES SNAP P_REAL OUTP INIT |
| """ |
| import sys, time |
| import torch, torch.nn.functional as F |
| from torch.utils.data import DataLoader |
| KIND, MINUTES, SNAP, P_REAL, OUTP, INIT = sys.argv[1], float(sys.argv[2]), float(sys.argv[3]), float(sys.argv[4]), sys.argv[5], sys.argv[6] |
| sys.argv = sys.argv[:1] + ["60", "32"] |
| from jl_synth import Synth5 |
| from train3 import dice |
| init = torch.load(INIT, map_location="cpu", weights_only=True)["state_dict"] |
| if KIND == "siam": |
| from model_v2 import SiamCD |
| model = SiamCD(); model.load_state_dict(init); model = model.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": 2e-5}, {"params": rest, "base": 1e-4}], lr=1e-4, weight_decay=1e-4) |
| else: |
| import segmentation_models_pytorch as smp |
| model = smp.Unet(encoder_name="resnet18", encoder_weights=None, in_channels=6, classes=3); model.load_state_dict(init); 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() |
| dl = DataLoader(Synth5(10**7, P_REAL), batch_size=32 if KIND == "siam" else 64, num_workers=18, 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) |
| for g in opt.param_groups: g["lr"] = g["base"] * min(1, (step + 1) / 100) |
| if KIND == "siam": |
| 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)) |
| loss = torch.nan_to_num(loss) |
| else: |
| y = torch.zeros_like(sp); y[lab[:, 1] > 0] = 2; y[lab[:, 0] > 0] = 1 |
| x = torch.cat([(pre - MEAN) / STD, (post - MEAN) / STD], 1) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| out = model(x) |
| 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 |
| 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 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 el > MINUTES * 60: break |
| print("TRAIN_DONE", step, flush=True) |
|
|
