图像、两视图与同源正例 · 实现代码

def make_data(seed=SEED):
    rng=np.random.default_rng(seed)
    yy,xx=np.mgrid[:16,:16]
    images=[]; shape=[]; color=[]; split=[]
    # Stratified before view construction. One independently generated source per ID.
    for s in range(3):
        for c in range(2):
            for i in range(140):
                cx,cy=rng.uniform(6,10,2); r=rng.uniform(3.0,4.7)
                dx,dy=xx-cx,yy-cy
                if s==0: mask=dx*dx+dy*dy<=r*r
                elif s==1: mask=np.maximum(abs(dx),abs(dy))<=r
                else: mask=abs(dx)+abs(dy)<=r*1.35
                # Red and blue have equal channel sums: grayscale erases hue by design.
                rgb=np.array([.9,.15,.15] if c==0 else [.15,.15,.9])
                bg=rng.uniform(.015,.075,(16,16,1))
                im=np.where(mask[...,None],rgb*rng.uniform(.8,1),bg)
                im+=rng.normal(0,.012,im.shape)
                images.append(np.clip(im,0,1).transpose(2,0,1))
                shape.append(s); color.append(c); split.append(0 if i<100 else 1)
    x=np.asarray(images,dtype='float32'); split=np.asarray(split)
    ids=np.arange(len(x)); train=ids[split==0]; test=ids[split==1]
    assert set(train).isdisjoint(test)
    assert len({hashlib.sha256(v.tobytes()).hexdigest() for v in x})==len(x)
    return torch.from_numpy(x),np.asarray(shape),np.asarray(color),train,test

def views(x, destroy_color=False):
    # Independent transforms of each source; no labels/scene attributes are read.
    a=x.clone(); n=len(a)
    shifts=torch.randint(-1,2,(n,2))
    for i,(dy,dx) in enumerate(shifts.tolist()):
        a[i]=torch.roll(a[i],(dy,dx),(-2,-1))
    a=a*torch.empty(n,1,1,1).uniform_(.8,1.2)
    a=a+torch.randn_like(a)*.025
    # The ablation changes only the hue swap. Every run consumes this RNG draw.
    swap=torch.rand(n)<.5
    if destroy_color: a[swap]=a[swap][:,[2,1,0]]
    return a.clamp(0,1)

class ImageEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv=nn.Sequential(nn.Conv2d(3,12,3,padding=1),nn.ReLU(),
          nn.Conv2d(12,24,3,stride=2,padding=1),nn.ReLU(),
          nn.Conv2d(24,24,3,stride=2,padding=1),nn.ReLU())
        self.readout=nn.Linear(24*4*4,48)
    def forward(self,x): return self.readout(self.conv(x).flatten(1))

class ContrastiveModel(nn.Module):
    def __init__(self):
        super().__init__(); self.encoder=ImageEncoder()
        self.projector=nn.Sequential(nn.Linear(48,48),nn.ReLU(),nn.Linear(48,24))
    def forward(self,x):
        h=self.encoder(x); z=self.projector(h)
        return h,z

def info_nce(z1,z2,tau=.2):
    z=F.normalize(torch.cat([z1,z2]),dim=1); n=len(z1)
    logits=z@z.T/tau
    logits=logits.masked_fill(torch.eye(2*n,dtype=torch.bool),-1e9)
    positive=(torch.arange(2*n)+n)%(2*n)
    return F.cross_entropy(logits,positive)

def train_simclr(initial,x,train,destroy_color,epochs=40):
    # Equal initialization, minibatch ordering, optimizer, steps, and RNG streams.
    torch.manual_seed(SEED+1)
    model=copy.deepcopy(initial); opt=torch.optim.Adam(model.parameters(),lr=.002)
    history=[]; before=next(model.encoder.parameters()).detach().clone()
    for epoch in range(epochs):
        model.train(); losses=[]
        for ids in torch.as_tensor(train)[torch.randperm(len(train))].split(60):
            batch=x[ids]; a,b=views(batch,destroy_color),views(batch,destroy_color)
            loss=info_nce(model(a)[1],model(b)[1])
            opt.zero_grad(); loss.backward(); opt.step(); losses.append(float(loss.detach()))
        history.append(float(np.mean(losses)))
    delta=float((next(model.encoder.parameters()).detach()-before).norm())
    assert delta>0
    return model,history,delta