Skip to content

CD-k on Two Moons

intermediate · Train an energy network with Contrastive Divergence and Langevin negatives.

Run it:

1
python examples/20-training/01-mcmc-losses/01-cd-k/main.py

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
"""Contrastive Divergence: train an energy net so data sits in low-energy regions.

CD pushes E down on real data and up on short-run MCMC "negatives" from the model.
"""

import os

import torch
from torch import nn

from torchebm.core import BaseModel
from torchebm.datasets import TwoMoonsDataset
from torchebm.losses import ContrastiveDivergence
from torchebm.samplers import LangevinDynamics

SMOKE = os.getenv("TORCHEBM_SMOKE") == "1"
N_STEPS = 20 if SMOKE else 1000


class MLPEnergy(BaseModel):           # E(x): R^2 -> R, one scalar energy per point
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 128), nn.SiLU(),
            nn.Linear(128, 128), nn.SiLU(),
            nn.Linear(128, 1),
        )

    def forward(self, x):
        return self.net(x).squeeze(-1)


data = TwoMoonsDataset(n_samples=3000, noise=0.05, seed=0).get_data()

energy = MLPEnergy()
sampler = LangevinDynamics(model=energy, step_size=0.1, noise_scale=1.0)  # draws negatives
cd = ContrastiveDivergence(model=energy, sampler=sampler, k_steps=10)     # CD-10
opt = torch.optim.Adam(energy.parameters(), lr=1e-3)

for step in range(N_STEPS):
    batch = data[torch.randint(len(data), (256,))]
    loss, _negatives = cd(batch)      # gradient of loss shapes the energy landscape
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 200 == 0:
        print(f"step {step:4d}  loss {loss.item():+.3f}")