`visual-search/train/finetune.py`

101 lines
#!/usr/bin/env python3
# Fine-tune open_clip ViT-B-32 (laion2b) on the DW catalog so image embeddings specialize to
# wallcovering/fabric patterns → sharper "really ID the pattern" retrieval than off-the-shelf CLIP.
# Same architecture + pretrained tag as the live search service, so the fine-tuned visual weights
# are a DROP-IN replacement. Standard symmetric CLIP contrastive (InfoNCE) loss. Runs on Apple MPS.
#
#   PYTORCH_ENABLE_MPS_FALLBACK=1 ./venv/bin/python finetune.py --epochs 4 --batch 96
import os, sys, json, time, argparse, random
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
import open_clip

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

def clean(s): return ' '.join((s or '').replace('|', ' ').split())[:120]

def caption(pattern, color, vendor, ptype):
    # natural DW-domain caption: "<pattern> in <color>, <vendor> <type>"
    bits = [clean(pattern)]
    if color and color.strip(): bits.append('in ' + clean(color))
    tail = ' '.join(x for x in [clean(vendor), clean(ptype).lower()] if x)
    txt = ', '.join([b for b in bits if b])
    if tail: txt += ', ' + tail
    return txt.lower().strip(', ')

class DWSet(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]
        img = self.pp(Image.open(os.path.join(CACHE, idv + '.jpg')).convert('RGB'))
        txt = self.tok([caption(pat, col, ven, typ)])[0]
        return img, txt

def load_rows():
    rows = []
    with open(os.path.join(HERE, 'manifest.tsv')) as f:
        for line in f:
            p = line.rstrip('\n').split('\t')
            if len(p) < 6: continue
            idv = p[0]
            if os.path.exists(os.path.join(CACHE, idv + '.jpg')):
                rows.append((idv, p[2], p[3], p[4], p[5]))   # id, pattern, color, vendor, type
    return rows

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--epochs', type=int, default=4)
    ap.add_argument('--batch', type=int, default=96)
    ap.add_argument('--lr', type=float, default=3e-6)      # gentle — avoid catastrophic forgetting
    ap.add_argument('--max_hours', type=float, default=7.5)
    a = ap.parse_args()

    dev = 'mps' if torch.backends.mps.is_available() else 'cpu'
    print(f"device={dev}", flush=True)
    model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
    tokenizer = open_clip.get_tokenizer('ViT-B-32')
    model = model.to(dev).float(); model.train()

    rows = load_rows()
    if len(rows) < 64:
        print(f"only {len(rows)} cached images — need more; is download.py done? exiting", flush=True); return
    random.seed(0); random.shuffle(rows)
    print(f"training pairs (cached): {len(rows)}", flush=True)
    ds = DWSet(rows, preprocess, tokenizer)
    dl = DataLoader(ds, batch_size=a.batch, shuffle=True, num_workers=4, 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:
                el = (time.time() - t0) / 3600
                print(f"ep{ep} step{step} loss={loss.item():.4f} elapsed={el:.2f}h", flush=True)
            if (time.time() - t0) / 3600 > a.max_hours:
                print("max_hours reached — stopping", flush=True)
                torch.save(model.state_dict(), os.path.join(CKPT, 'dw_clip_ft.pt')); return
        ckpt = os.path.join(CKPT, f'dw_clip_ft_ep{ep}.pt')
        torch.save(model.state_dict(), ckpt)
        print(f"saved {ckpt}", flush=True)
    torch.save(model.state_dict(), os.path.join(CKPT, 'dw_clip_ft.pt'))
    print(f"DONE — final weights → ckpt/dw_clip_ft.pt  ({(time.time()-t0)/3600:.2f}h)", flush=True)

if __name__ == '__main__':
    main()