h 与 z：目标作用在哪一层？ · 实现代码

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