`visual-search/train/finetune_aug.py`

95 lines
#!/usr/bin/env python3
# AUGMENTATION re-fine-tune: warm-start from dw_clip_ft.pt and continue training with heavy
# handheld-photo augmentation (rotate/crop/blur/glare/JPEG) on the IMAGE side, so the model learns
# that a degraded phone shot maps to the same catalog identity — the gap robustness_test.py exposed
# (28% top-1 on degraded photos). Caption stays clean; only the image is degraded during training.
#   PYTORCH_ENABLE_MPS_FALLBACK=1 ./venv/bin/python finetune_aug.py --epochs 3 --batch 96
import os, sys, io, time, random, argparse
os.environ.setdefault('PYTORCH_ENABLE_MPS_FALLBACK', '1')
import torch, torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from PIL import Image, ImageFilter, ImageEnhance
import open_clip
import finetune as T   # reuse caption() + load_rows()

HERE = os.path.dirname(os.path.abspath(__file__))
CACHE = os.path.join(HERE, 'cache'); CKPT = os.path.join(HERE, 'ckpt')
Image.MAX_IMAGE_PIXELS = None

def augment(im):
    # simulate a handheld iPhone shot of a physical sample — varied strength, each applied ~probabilistically
    if im.width < 48 or im.height < 48: return im       # too small to degrade safely (corrupt/tiny cache)
    w, h = im.size
    if random.random() < 0.7: im = im.rotate(random.uniform(-8, 8), expand=False, fillcolor=(238,238,238))
    if random.random() < 0.7:
        cx, cy = random.uniform(0.02, 0.10), random.uniform(0.02, 0.10)
        l,t,r,bt = int(w*cx), int(h*cy), int(w*(1-cx)), int(h*(1-cy))
        if r-l >= 32 and bt-t >= 32: im = im.crop((l,t,r,bt))   # never crop to an empty/tiny box
    if random.random() < 0.6: im = im.filter(ImageFilter.GaussianBlur(random.uniform(0.4, 1.6)))
    if random.random() < 0.7: im = ImageEnhance.Brightness(im).enhance(random.uniform(0.85, 1.30))   # glare / shadow
    if random.random() < 0.5: im = ImageEnhance.Contrast(im).enhance(random.uniform(0.80, 1.15))
    if random.random() < 0.6 and im.width >= 32 and im.height >= 32:
        try:
            b = io.BytesIO(); im.save(b, 'JPEG', quality=random.randint(45, 80)); b.seek(0); im = Image.open(b).convert('RGB')
        except Exception: pass                          # a bad recompress must never crash training
    return im

class AugSet(Dataset):
    def __init__(self, rows, preprocess, tokenizer):
        self.rows = rows; self.pp = preprocess; self.tok = tokenizer
    def __len__(self): return len(self.rows)
    def __getitem__(self, i):
        idv, pat, col, ven, typ = self.rows[i]
        try:
            im = Image.open(os.path.join(CACHE, idv + '.jpg')).convert('RGB')
            img = self.pp(augment(im))                   # <-- degrade before CLIP preprocess
        except Exception:
            return self.__getitem__((i + 1) % len(self.rows))   # corrupt cache image → use a neighbor, never crash
        txt = self.tok([T.caption(pat, col, ven, typ)])[0]
        return img, txt

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--epochs', type=int, default=3)
    ap.add_argument('--batch', type=int, default=96)
    ap.add_argument('--lr', type=float, default=2e-6)      # gentle — warm-starting an already-good model
    ap.add_argument('--max_hours', type=float, default=3.0)
    a = ap.parse_args()

    dev = 'mps' if torch.backends.mps.is_available() else 'cpu'
    model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
    tokenizer = open_clip.get_tokenizer('ViT-B-32')
    warm = os.path.join(CKPT, 'dw_clip_ft.pt')             # WARM-START from the current fine-tuned model
    model.load_state_dict(torch.load(warm, map_location='cpu'), strict=False)
    print('warm-started from dw_clip_ft.pt · device', dev, flush=True)
    model = model.to(dev).float(); model.train()

    rows = T.load_rows()
    print(f'training pairs (cached): {len(rows)}', flush=True)
    dl = DataLoader(AugSet(rows, preprocess, tokenizer), batch_size=a.batch, shuffle=True,
                    num_workers=5, drop_last=True, persistent_workers=True)
    opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=0.1)
    t0 = time.time(); step = 0
    for ep in range(a.epochs):
        for imgs, txts in dl:
            imgs = imgs.to(dev); txts = txts.to(dev)
            imf = F.normalize(model.encode_image(imgs), dim=-1)
            txf = F.normalize(model.encode_text(txts), dim=-1)
            scale = model.logit_scale.exp().clamp(max=100)
            logits = scale * imf @ txf.t()
            labels = torch.arange(logits.size(0), device=dev)
            loss = (F.cross_entropy(logits, labels) + F.cross_entropy(logits.t(), labels)) / 2
            opt.zero_grad(); loss.backward(); opt.step()
            with torch.no_grad(): model.logit_scale.clamp_(0, 4.6052)
            step += 1
            if step % 20 == 0: print(f'ep{ep} step{step} loss={loss.item():.4f} elapsed={(time.time()-t0)/3600:.2f}h', flush=True)
            if (time.time()-t0)/3600 > a.max_hours:
                torch.save(model.state_dict(), os.path.join(CKPT, 'dw_clip_ft_aug.pt')); print('max_hours — saved', flush=True); return
        torch.save(model.state_dict(), os.path.join(CKPT, f'dw_clip_ft_aug_ep{ep}.pt'))
        print(f'saved ep{ep}', flush=True)
    torch.save(model.state_dict(), os.path.join(CKPT, 'dw_clip_ft_aug.pt'))
    print(f'DONE — aug weights → ckpt/dw_clip_ft_aug.pt ({(time.time()-t0)/3600:.2f}h)', flush=True)

if __name__ == '__main__':
    main()