def sample_data(nsample=100, goal=4):
x1 = np.linspace(start=-5, stop=5, num=100)
x2 = np.linspace(start=-5, stop=5, num=100)
samples, samples_ind = [], []
weights = []
while len(samples) < nsample:
x1_rand = np.random.choice(range(len(x1)))
x2_rand = np.random.choice(range(len(x2)))
sample = np.array([x1[x1_rand], x2[x2_rand]])
fval = ackley(sample)
if fval < goal:
weights.append(1)
else:
weights.append(-1)
samples_ind.append(np.array([x1_rand, x2_rand]))
samples.append(sample)
samples_arr = np.stack(samples)
samples_ind_arr = np.stack(samples_ind)
data = np.stack([samples_arr, samples_ind_arr], axis=1)
weights = np.array(weights)
return data, weights
def collect_samples(self, n_samples, uniform=False):
x1, x2 = self.input_vectors
n_collected = 0
new_samples = []
while n_collected < n_samples:
if uniform:
_, xnew_id_np = self.sample_data(1, self.goal)
xnew_id_np = xnew_id_np.astype('int')
else:
_, xnew_ind = self.sample_model(1)
xnew_id_np = xnew_ind.to(self.cpu).data.numpy().astype('int')
xsample = np.array([[x1[xnew_id_np[0, 0]],
x2[xnew_id_np[0, 1]]]]).astype('float32')
# simulate and compute the adjustment weights
fval = ackley(xsample * 5)
if xsample not in self.buffer:
self.buffer.add_samples(xsample, xnew_id_np, fval)
new_samples.append(xsample)
n_collected += 1
else:
print(f'item {xsample} already exists!')
return new_samples
def sample_data(self, nsample=100, goal=4):
x1, x2 = self.input_vectors
samples_ind = []
weights = []
while len(samples_ind) < nsample:
x1_rand = np.random.choice(range(len(x1)))
x2_rand = np.random.choice(range(len(x2)))
sample = np.array([x1[x1_rand], x2[x2_rand]])
fval = ackley(sample)
if fval < goal:
weights.append(1)
else:
weights.append(-1)
samples_ind.append(np.array([x1_rand, x2_rand]))
data_ind = np.stack(samples_ind)
samples = np.stack([x1[data_ind[:, 0]], x2[data_ind[:, 1]]], axis=-1)
samples_norm, delta = self.normalize(samples)
data = np.stack([samples_norm, data_ind], axis=1)
weights = np.array(weights)
return data, delta, weights
def collect_samples(self, n_samples):
n_collected = 0
while n_collected < n_samples:
xnew, xnew_ind, xnew_probs = self.sample_model(
1, *self.input_vectors)
xnew_np = xnew.to(self.cpu).data.numpy()
xnew_id_np = xnew_ind.to(self.cpu).data.numpy()
xnew_probs_np = xnew_probs.to(self.cpu).data.numpy()
# simulate and compute the adjustment weights
fval = ackley(xnew_np)
if xnew_np not in self.buffer:
self.buffer.add_samples(xnew_np, xnew_id_np, fval)
n_collected += 1
else:
print(f'item {xnew_np} already exists!')
def train(self, n_samples, *input_vecs, iter_cnt: int):
xnew, xnew_ind, xnew_probs = self.sample_model(n_samples, *input_vecs)
xnew_np = xnew.to(self.cpu).data.numpy()
xnew_id_np = xnew_ind.to(self.cpu).data.numpy()
xnew_probs_np = xnew_probs.to(self.cpu).data.numpy()
# simulate and compute the adjustment weights
fval = ackley(xnew_np)
self.buffer.add_samples(xnew_np, xnew_id_np, fval, xnew_probs_np)
# treat the sampled data as a static data set and take some gradient steps on it
xtr, xte, wtr, wte = self.buffer.draw_tr_te_ds()
# plotting
self.viz(xtr[:, 0, :], 'training', f'{iter_cnt}_before')
# per epoch
print('-' * 50)
for epoch_id in range(self.nepochs):
tr_nll = self.run_epoch(xtr, wtr, mode='train')
te_nll = self.run_epoch(xte, wte, mode='test')
print(f'[train_{iter_cnt}] epoch {epoch_id} loss = {tr_nll}')
print(f'[test_{iter_cnt}] epoch {epoch_id} loss = {te_nll}')
def main(self, seed=10):
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
data = sample_data(self.nsample)
fval = ackley(data[:, 0, :])
weights = weight(fval, self.goal, 4, mode='le')
xtr, xte, wtr, wte = split_data(data, label=weights)
D = self.dim
self.model: nn.Module = MADE(D, self.hiddens, D * 100, seed=seed)
self.model.to(self.device)
self.opt = torch.optim.Adam(self.model.parameters(),
self.lr,
weight_decay=0)
B = self.bsize
N, D, _ = xtr.shape
# per epoch
tr_nll, te_nll = [], []
for epoch_id in range(self.nepoch):
nstep = N // B
# per batch
tr_nll_per_b, te_nll_per_b = 0, 0
for step in range(nstep):
self.model.train()
xb = xtr[step * B:step * B + B]
wb = wtr[step * B:step * B + B]
xb_tens = torch.from_numpy(xb).to(self.device)
wb_tens = torch.from_numpy(wb).to(self.device)
xin = xb_tens[:, 0, :]
xin_ind = xb_tens[:, 1, :].long()
loss = self.get_nll(xin, xin_ind, weights=wb_tens)
self.opt.zero_grad()
loss.backward()
self.opt.step()
# print(loss)
# for name, param in self.model.named_parameters():
# print(f'{name} = {param.grad}')
# import pdb
# pdb.set_trace()
tr_nll_per_b += loss.to(self.cpu).item() / nstep
self.model.eval()
xte_tens = torch.from_numpy(xte).to(self.device)
wte_tens = torch.from_numpy(wte).to(self.device)
xin_te = xte_tens[:, 0, :]
xin_ind_te = xte_tens[:, 1, :].long()
te_loss = self.get_nll(xin_te, xin_ind_te, weights=wte_tens)
te_nll.append(te_loss)
print(f'epoch = {epoch_id}, tr_nll = {tr_nll_per_b}')
print(f'epoch = {epoch_id}, te_nll = {te_loss}')
tr_nll.append(tr_nll_per_b)
self.plot_learning(tr_nll, te_nll)
x1 = np.linspace(start=-5, stop=5, num=100)
x2 = np.linspace(start=-5, stop=5, num=100)
samples, _ = self.sample_model(self.nsample, x1, x2)
samples = samples.to(self.cpu).data.numpy()
plot_data(samples,
scatter_loc='figs/test_model3_scatter.png',
hist_loc='figs/test_model3_hist2D.png')