def test(hparams):
# load dataset
N_test = 5000
model = torch.load(hparams.model_path)
x_start = torch.randn([N_test, 2])
x_end = sample_langevin(x_start, model, hparams.stepsize, hparams.n_steps)
plot_figure(x_start, x_end, save_path='Figures/ebm.pdf')
def train(hparams):
#wandb.init(project="ebm-gaussians")
seed_everything(hparams.seed)
model = mlp(sizes=[2, 100, 100, 1], activation=nn.ReLU)
optimizer = Adam(model.parameters(), lr=hparams.lr)
# load dataset
N_train = 5000
X_train = sample_data(N_train)
train_dl = DataLoader(X_train, batch_size=100, shuffle=True, num_workers=8)
losses = []
for _ in range(hparams.n_epochs):
for x in train_dl:
neg_x = torch.randn_like(x)
neg_x = sample_langevin(neg_x, model, hparams.stepsize,
hparams.n_steps)
optimizer.zero_grad()
pos_out = model(x)
neg_out = model(neg_x)
loss = (pos_out -
neg_out) + hparams.alpha * (pos_out**2 + neg_out**2)
loss = loss.mean()
loss.backward()
#torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.1)
optimizer.step()
losses.append(loss.item())
# wandb.log({'loss': loss.item()})
print('saving a trained model')
torch.save(model, hparams.model_path)
def test_langevin():
X = torch.tensor([[1,2],[4,5],[7,8],[11,12],[14,15]], dtype = torch.float)
model = FCNet(2, 1, l_hidden=(50,))
sample = sample_langevin(X, model, 0.1, 1)
model = FCNet(in_dim=2, out_dim=1, l_hidden=(100, 100), activation='relu', out_activation='linear')
#model.cuda()
#opt = Adam(model.parameters(), lr=args.lr)
opt = SGLD(model.parameters(), lr=args.lr)
# train loop
for i_epoch in range(args.n_epoch):
l_loss = []
for pos_x, in train_dl:
#pos_x = pos_x.cuda()
neg_x = torch.randn_like(pos_x)
neg_x = sample_langevin(neg_x, model, args.stepsize, args.n_step, intermediate_samples=False)
opt.zero_grad()
pos_out = model(pos_x)
neg_out = model(neg_x)
loss = (pos_out - neg_out) + args.alpha * (pos_out ** 2 + neg_out ** 2)
loss = loss.mean()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.1)
opt.step()
l_loss.append(loss.item())
print(np.mean(l_loss))
def test_langevin():
X = torch.randn(100, 2)
model = FCNet(2, 1, l_hidden=(50, ))
sample = sample_langevin(X, model, 0.1, 10)
# prepare for contour plot
energy_fn = potential_fn(args.energy_function)
xs = np.linspace(-grid_lim, grid_lim, n_grid)
ys = np.linspace(-grid_lim, grid_lim, n_grid)
XX, YY = np.meshgrid(xs, ys)
grids = np.stack([XX.flatten(), YY.flatten()]).T
e_grid = energy_fn(torch.tensor(grids))
# run langevin dynamics
grad_log_p = lambda x: -energy_fn(x)
x0 = torch.randn(n_grid, 2)
l_sample, l_dynamics = sample_langevin(x0,
grad_log_p,
stepsize,
n_steps,
intermediate_samples=True)
# plot
fig = plt.figure()
ax = plt.gca()
plt.axis('equal')
point = plt.scatter([], [])
if args.no_arrow:
arrow = None
else:
arrow = plt.quiver([0], [0], [1], [1],
scale=0.5,
scale_units='xy',