def main(args):
# load data
print('loading training data...')
dataset_directory = get_data_directory(__file__)
dataset_path = os.path.join(dataset_directory, 'faces_training.csv')
if not os.path.exists(dataset_path):
try:
os.makedirs(dataset_directory)
except OSError as e:
if e.errno != errno.EEXIST:
raise
pass
wget.download(
'https://d2hg8soec8ck9v.cloudfront.net/datasets/faces_training.csv',
dataset_path)
data = torch.tensor(np.loadtxt(dataset_path, delimiter=',')).float()
sparse_gamma_def = SparseGammaDEF()
# Due to the special logic in the custom guide (e.g. parameter clipping), the custom guide
# seems to be more amenable to higher learning rates.
# Nevertheless, the easy guide performs the best (presumably because of numerical instabilities
# related to the gamma distribution in the custom guide).
learning_rate = 0.2 if args.guide in ['auto', 'easy'] else 4.5
momentum = 0.05 if args.guide in ['auto', 'easy'] else 0.1
opt = optim.AdagradRMSProp({"eta": learning_rate, "t": momentum})
# use one of our three different guide types
if args.guide == 'auto':
guide = AutoDiagonalNormal(sparse_gamma_def.model,
init_loc_fn=init_to_feasible)
elif args.guide == 'easy':
guide = MyEasyGuide(sparse_gamma_def.model)
else:
guide = sparse_gamma_def.guide
# this is the svi object we use during training; we use TraceMeanField_ELBO to
# get analytic KL divergences
svi = SVI(sparse_gamma_def.model, guide, opt, loss=TraceMeanField_ELBO())
# we use svi_eval during evaluation; since we took care to write down our model in
# a fully vectorized way, this computation can be done efficiently with large tensor ops
svi_eval = SVI(sparse_gamma_def.model,
guide,
opt,
loss=TraceMeanField_ELBO(num_particles=args.eval_particles,
vectorize_particles=True))
print('\nbeginning training with %s guide...' % args.guide)
# the training loop
for k in range(args.num_epochs):
loss = svi.step(data)
# for the custom guide we clip parameters after each gradient step
if args.guide == 'custom':
clip_params()
if k % args.eval_frequency == 0 and k > 0 or k == args.num_epochs - 1:
loss = svi_eval.evaluate_loss(data)
print("[epoch %04d] training elbo: %.4g" % (k, -loss))
def __init__(self,
net,
prior,
guide_builder=None,
name="",
closed_form_kl=True):
super().__init__(net, prior, guide_builder=guide_builder, name=name)
self.cached_output = None
self.cached_kl_loss = None
self._loss = TraceMeanField_ELBO() if closed_form_kl else Trace_ELBO()
def main(args):
# load data
print('loading training data...')
dataset_directory = get_data_directory(__file__)
dataset_path = os.path.join(dataset_directory, 'faces_training.csv')
if not os.path.exists(dataset_path):
try:
os.makedirs(dataset_directory)
except OSError as e:
if e.errno != errno.EEXIST:
raise
pass
wget.download('https://d2fefpcigoriu7.cloudfront.net/datasets/faces_training.csv', dataset_path)
data = torch.tensor(np.loadtxt(dataset_path, delimiter=',')).float()
sparse_gamma_def = SparseGammaDEF()
# due to the special logic in the custom guide (e.g. parameter clipping), the custom guide
# is more numerically stable and enables us to use a larger learning rate (and consequently
# achieves better results)
learning_rate = 0.2 if args.auto_guide else 4.5
momentum = 0.05 if args.auto_guide else 0.1
opt = optim.AdagradRMSProp({"eta": learning_rate, "t": momentum})
# either use an automatically constructed guide (see pyro.contrib.autoguide for details) or our custom guide
guide = AutoDiagonalNormal(sparse_gamma_def.model) if args.auto_guide else sparse_gamma_def.guide
# this is the svi object we use during training; we use TraceMeanField_ELBO to
# get analytic KL divergences
svi = SVI(sparse_gamma_def.model, guide, opt, loss=TraceMeanField_ELBO())
# we use svi_eval during evaluation; since we took care to write down our model in
# a fully vectorized way, this computation can be done efficiently with large tensor ops
svi_eval = SVI(sparse_gamma_def.model, guide, opt,
loss=TraceMeanField_ELBO(num_particles=args.eval_particles, vectorize_particles=True))
guide_description = 'automatically constructed' if args.auto_guide else 'custom'
print('\nbeginning training with %s guide...' % guide_description)
# the training loop
for k in range(args.num_epochs):
loss = svi.step(data)
if not args.auto_guide:
# for the custom guide we clip parameters after each gradient step
sparse_gamma_def.clip_params()
if k % args.eval_frequency == 0 and k > 0 or k == args.num_epochs - 1:
loss = svi_eval.evaluate_loss(data)
print("[epoch %04d] training elbo: %.4g" % (k, -loss))
def train(gpmodule, optimizer=None, loss_fn=None, retain_graph=None, num_steps=1000):
"""
A helper to optimize parameters for a GP module.
:param ~pyro.contrib.gp.models.GPModel gpmodule: A GP module.
:param ~torch.optim.Optimizer optimizer: A PyTorch optimizer instance.
By default, we use Adam with ``lr=0.01``.
:param callable loss_fn: A loss function which takes inputs are
``gpmodule.model``, ``gpmodule.guide``, and returns ELBO loss.
By default, ``loss_fn=TraceMeanField_ELBO().differentiable_loss``.
:param bool retain_graph: An optional flag of ``torch.autograd.backward``.
:param int num_steps: Number of steps to run SVI.
:returns: a list of losses during the training procedure
:rtype: list
"""
optimizer = (
torch.optim.Adam(gpmodule.parameters(), lr=0.01)
if optimizer is None
else optimizer
)
# TODO: add support for JIT loss
loss_fn = TraceMeanField_ELBO().differentiable_loss if loss_fn is None else loss_fn
def closure():
optimizer.zero_grad()
loss = loss_fn(gpmodule.model, gpmodule.guide)
torch_backward(loss, retain_graph)
return loss
losses = []
for i in range(num_steps):
loss = optimizer.step(closure)
losses.append(torch_item(loss))
return losses
def getELBO(self, elbo_type, particles):
if elbo_type == 'TraceELBO':
return Trace_ELBO(num_particles=particles)
elif elbo_type == "MeanFieldELBO":
return TraceMeanField_ELBO(num_particles=particles)
else:
raise ValueError("{} ELBO not supported".format(elbo_type))
def main(args):
# load data
print('loading training data...')
dataset_directory = get_data_directory(__file__)
dataset_path = os.path.join(dataset_directory, 'faces_training.csv')
if not os.path.exists(dataset_path):
try:
os.makedirs(dataset_directory)
except OSError as e:
if e.errno != errno.EEXIST:
raise
pass
wget.download(
'https://d2fefpcigoriu7.cloudfront.net/datasets/faces_training.csv',
dataset_path)
data = torch.tensor(np.loadtxt(dataset_path, delimiter=',')).float()
learning_rate = 4.5
momentum = 0.1
opt = optim.AdagradRMSProp({"eta": learning_rate, "t": momentum})
# this is the svi object we use during training; we use TraceMeanField_ELBO to
# get analytic KL divergences
svi = SVI(model, guide, opt, loss=TraceMeanField_ELBO())
# we use svi_eval during evaluation; since we took care to write down our model in
# a fully vectorized way, this computation can be done efficiently with large tensor ops
svi_eval = SVI(model_original,
guide,
opt,
loss=TraceMeanField_ELBO(num_particles=args.eval_particles,
vectorize_particles=True))
guide_description = 'custom'
print('\nbeginning training with %s guide...' % guide_description)
# the training loop
for k in range(args.num_epochs):
loss = svi.step(data)
clip_params()
if k % args.eval_frequency == 0 and k > 0 or k == args.num_epochs - 1:
loss = svi_eval.evaluate_loss(data)
print("[epoch %04d] training elbo: %.4g" % (k, -loss))
def _train_svi(self, train_loader, epochs, lr, device):
self.device = device
print("\n == SVI training ==")
optimizer = pyro.optim.Adam({"lr": lr})
elbo = TraceMeanField_ELBO()
svi = SVI(self.model, self.guide, optimizer, loss=elbo)
loss_list = []
accuracy_list = []
start = time.time()
for epoch in range(epochs):
loss = 0.0
correct_predictions = 0.0
for x_batch, y_batch in train_loader:
x_batch = x_batch.to(device)
y_batch = y_batch.to(device)
labels = y_batch.argmax(-1)
loss += svi.step(x_data=x_batch, y_data=labels)
outputs = self.forward(x_batch, n_samples=10)
predictions = outputs.argmax(dim=-1)
correct_predictions += (predictions == labels).sum().item()
if DEBUG:
print("\n",
pyro.get_param_store()["model.0.weight_loc"][0][:5])
print("\n", predictions[:10], "\n", labels[:10])
total_loss = loss / len(train_loader.dataset)
accuracy = 100 * correct_predictions / len(train_loader.dataset)
print(
f"\n[Epoch {epoch + 1}]\t loss: {total_loss:.2f} \t accuracy: {accuracy:.2f}",
end="\t")
loss_list.append(loss)
accuracy_list.append(accuracy)
execution_time(start=start, end=time.time())
self.save()
plot_loss_accuracy(dict={
'loss': loss_list,
'accuracy': accuracy_list
},
path=TESTS + self.name + "/" + self.name +
"_training.png")
def train(self,
*,
raw_expr,
encoded_expr,
num_epochs=100,
batch_size=32,
learning_rate=1e-3,
posterior_samples=20):
logging.info('Validating data ...')
assert (raw_expr.shape == encoded_expr.shape)
read_depth = torch.tensor(raw_expr.sum(-1)[:,
np.newaxis]).to(self.device)
raw_expr = torch.tensor(raw_expr).to(self.device)
encoded_expr = torch.tensor(encoded_expr).to(self.device)
logging.info('Initializing model ...')
pyro.clear_param_store()
adam_params = {"lr": 1e-3}
optimizer = Adam(adam_params)
svi = SVI(self.model,
self.guide,
optimizer,
loss=TraceMeanField_ELBO())
logging.info('Training ...')
num_batches = read_depth.shape[0] // batch_size
bar = trange(num_epochs)
try:
for epoch in bar:
running_loss = 0.0
for batch in self.epoch_batch(raw_expr,
encoded_expr,
read_depth,
batch_size=batch_size):
loss = svi.step(*batch)
running_loss += loss / batch_size
bar.set_postfix(epoch_loss='{:.2e}'.format(running_loss))
except KeyboardInterrupt:
logging.error('Interrupted training.')
self.summarize_posterior(raw_expr,
encoded_expr,
read_depth,
num_samples=posterior_samples)
return self
def log_prob(self, window_begin, window_end, truth):
forecast_hours = len(truth)
self.args.funsor = False # sets behavior of model and guide
assert 0 <= window_begin < window_end < window_end + forecast_hours <= len(
self.counts)
features = self.features[window_begin: window_end + forecast_hours] \
.to(device=self.args.device)
x = self.counts[window_begin:window_end].to(device=self.args.device)
y = truth.to(device=self.args.device)
xy = torch.cat([x, y], dim=0)
loss = TraceMeanField_ELBO().loss
logp_x = -loss(self.model, self.guide, features[:len(x)], x)
logp_xy = -loss(self.model, self.guide, features[:len(xy)], xy)
return logp_xy - logp_x
def train(self,
*,
accessibility_matrix,
num_epochs=125,
batch_size=32,
learning_rate=1e-3,
eval_every=10,
test_proportion=0.05):
logging.info('Initializing model ...')
pyro.clear_param_store()
optimizer = Adam({"lr": 1e-3})
self.svi = SVI(self.model,
self.guide,
optimizer,
loss=TraceMeanField_ELBO())
test_set = np.random.rand(
accessibility_matrix.shape[0]) < test_proportion
train_set = ~test_set
logging.info("Training with {} cells, testing with {}.".format(
str(train_set.sum()), str(test_set.sum())))
num_batches = accessibility_matrix.shape[0] // batch_size
logging.info('Training for {} epochs'.format(str(num_epochs)))
try:
for epoch in range(1, num_epochs + 1):
running_loss = 0.0
for batch in self.epoch_batch(accessibility_matrix[train_set],
batch_size=batch_size):
loss = self.svi.step(*batch)
running_loss += loss / batch_size
logging.info('Done epoch {}/{}. Training loss: {:.3e}'.format(
str(epoch), str(num_epochs), running_loss))
if (epoch % eval_every == 0
or epoch == num_epochs) and test_set.sum() > 0:
test_logp = self.evaluate(accessibility_matrix[test_set],
batch_size=batch_size)
logging.info('Test logp: {:.4e}'.format(test_logp))
except KeyboardInterrupt:
logging.error('Interrupted training.')
self.summarize_posterior(accessibility_matrix)
return self
def fit(self,
data_loader,
optim,
num_epochs,
callback=None,
num_particles=1,
closed_form_kl=True,
device=None):
"""Optimizes the variational parameters on data from data_loader using optim for num_epochs.
:param Iterable data_loader: iterable over batches of data, e.g. a torch.utils.data.DataLoader. Assumes that
each element consists of a length two tuple of list, with the first element either containing a single
object or a list of objects, e.g. torch.Tensors, that are the inputs to the neural network. The second
element is a single torch.Tensor e.g. of class labels.
:param optim: pyro optimizer to be used for constructing an SVI object, e.g. pyro.optim.Adam({"lr": 1e-3}).
:param int num_epochs: number of passes over data_loader.
:param callable callback: optional function to invoke after every training epoch. Receives the BNN object,
the epoch number and the average value of the ELBO over the epoch. May return True to terminate
optimization before num_epochs, e.g. if it finds that a validation log likelihood saturates.
:param int num_particles: number of MC samples for estimating the ELBO.
:param bool closed_form_kl: whether to use TraceMeanField_ELBO or Trace_ELBO, i.e. calculate KL divergence
between approximate posterior and prior in closed form or via a Monte Carlo estimate.
:param torch.device device: optional device to send the data to.
"""
old_training_state = self.net.training
self.net.train(True)
loss = TraceMeanField_ELBO(
num_particles) if closed_form_kl else Trace_ELBO(num_particles)
svi = SVI(self.model, self.guide, optim, loss=loss)
for i in range(num_epochs):
elbo = 0.
num_batch = 1
for num_batch, (input_data, observation_data) in enumerate(
iter(data_loader), 1):
elbo += svi.step(tuple(_to(input_data, device)),
tuple(_to(observation_data, device))[0])
# the callback can stop training by returning True
if callback is not None and callback(self, i, elbo / num_batch):
break
self.net.train(old_training_state)
return svi
class PytorchBNN(GuidedBNN):
"""Low-level variational BNN class that can serve as a drop-in replacement for an nn.Module.
:param bool closed_form_kl: whether to use TraceMeanField_ELBO or Trace_ELBO as the loss, i.e. calculate KL
divergences in closed form or via a Monte Carlo approximate of the difference of log densities between
variational posterior and prior."""
def __init__(self,
net,
prior,
guide_builder=None,
name="",
closed_form_kl=True):
super().__init__(net, prior, guide_builder=guide_builder, name=name)
self.cached_output = None
self.cached_kl_loss = None
self._loss = TraceMeanField_ELBO() if closed_form_kl else Trace_ELBO()
def named_pytorch_parameters(self, *input_data):
"""Equivalent of the named_parameters method of an nn.Module. Ensures that prior and guide are run once to
initialize all pyro parameters. Those are then collected and returned via the trace poutine."""
model_trace = poutine.trace(self.net,
param_only=True).get_trace(*input_data)
guide_trace = poutine.trace(self.net_guide,
param_only=True).get_trace(*input_data)
for name, msg in itertools.chain(model_trace.nodes.items(),
guide_trace.nodes.items()):
yield name, msg["value"].unconstrained()
def pytorch_parameters(self, input_data_or_fwd_fn):
yield from map(itemgetter(1),
self.named_pytorch_parameters(input_data_or_fwd_fn))
def cached_forward(self, *args, **kwargs):
# cache the output of forward to make it effectful, so that we can access the output when running forward with
# posterior rather than prior samples
self.cached_output = super().forward(*args, **kwargs)
return self.cached_output
def forward(self, *args, **kwargs):
self.cached_kl_loss = self._loss.differentiable_loss(
self.cached_forward, self.net_guide, *args, **kwargs)
return self.cached_output
def train(self, data_loader, n_epochs, num_particles=1, lr=1e-3, log_per=5, show_smooth=True, save_per=10):
pyro.clear_param_store()
self.guide = partial(guide, model=self)
svi = SVI(self, self.guide, Adam({"lr": lr}), TraceMeanField_ELBO(num_particles=num_particles))
losses = []
fig = None
pp = ProgressPlotter(losses, "$-ELBO$", log_per, show_smooth)
pp.start()
for epoch in range(n_epochs):
total_loss = 0.
for x, y in data_loader:
x, y = x.float().to(self.device), y.float().to(self.device)
loss = svi.step(x, y) / y.numel()
total_loss += loss
total_loss /= len(data_loader)
losses.append(total_loss)
fig = pp.update(epoch)
if epoch % save_per == 1:
self.save('MultilayerBayesianTrick.pth')
return pp.fig, losses
def infer_parameters(self, loader, lr=0.01, momentum=0.9, num_epochs=30):
optim = pyroopt.SGD({'lr': lr, 'momentum': momentum, 'nesterov': True})
elbo = TraceMeanField_ELBO()
svi = SVI(self.model, self.guide, optim, elbo)
kl_factor = loader.batch_size / len(loader.dataset)
for i in range(num_epochs):
total_loss = 0.0
total = 0.0
correct = 0.0
for images, labels in loader:
loss = svi.step(images.cuda(),
labels.cuda(),
kl_factor=kl_factor)
pred = self.forward(images.cuda(), n_samples=1).mean(0)
total_loss += loss / len(loader.dataset)
total += labels.size(0)
correct += (pred.argmax(-1) == labels.cuda()).sum().item()
param_store = pyro.get_param_store()
print(
f"[Epoch {i + 1}] loss: {total_loss:.5E} accuracy: {correct / total * 100:.5f}"
)
def main():
assert pyro.__version__.startswith('1.6.0')
# Enable smoke test to test functionality
smoke_test = False
logging.info(f"CUDA available: {torch.cuda.is_available()}")
# Loading data
logging.info("Loading data...")
docs = prepro_file_load("doc_word_matrix").to_dense()
id2word = prepro_file_load("id2word")
# Put vocab into dataframe for exploration of data
vocab = pd.DataFrame(columns=['index', 'word'])
vocab['index'] = list(id2word.keys())
vocab['word'] = list(id2word.values())
logging.info(f"Dictionary size: {len(vocab)}")
logging.info(f"Corpus size: {docs.shape}")
# Setting global variables
seed = 0
torch.manual_seed(seed)
pyro.set_rng_seed(seed)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
docs = docs.float()
num_categories = 0
num_topics = 30 if not smoke_test else 3
batch_size = 32
learning_rate = 1e-3
num_epochs = 50 if not smoke_test else 1
# Training
pyro.clear_param_store()
prodLDA = ProdLDA(vocab_size=docs.shape[1],
num_topics=num_topics,
hidden=100 if not smoke_test else 10,
dropout=0.2)
prodLDA.to(device)
optimizer = pyro.optim.Adam({"lr": learning_rate})
svi = SVI(prodLDA.model,
prodLDA.guide,
optimizer,
loss=TraceMeanField_ELBO())
num_batches = int(math.ceil(docs.shape[0] /
batch_size)) if not smoke_test else 1
losses = []
logging.info("Training...")
bar = trange(num_epochs)
for epoch in bar:
running_loss = 0.0
for i in range(num_batches):
batch_docs = docs[i * batch_size:(i + 1) *
batch_size, :].to(device)
loss = svi.step(batch_docs)
running_loss += loss / batch_docs.size(0)
# Save and log losses
losses.append(running_loss)
bar.set_postfix(epoch_loss='{:.2e}'.format(running_loss))
if epoch % 5 == 0:
logging.info('{: >5d}\t{}'.format(epoch,
'{:.2e}'.format(running_loss)))
logging.info(f"Final loss: {'{:.2e}'.format(losses[-1])}/{losses[-1]}")
if not smoke_test:
# Plot loss over epochs
plt.plot(losses)
plt.title("ELBO")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plot_file_name = "../ProdLDA-loss-2017_categories-" + str(num_categories) + \
"_topics-" + str(num_topics) + \
"_batch-" + str(batch_size) + \
"_lr-" + str(learning_rate) + \
"_epochs-" + str(num_epochs) + \
".png"
plt.savefig(plot_file_name)
plt.show()
# Logging top 10 weighted words in topics
beta = prodLDA.beta()
for n in range(beta.shape[0]):
sorted_, indices = torch.sort(beta[n], descending=True)
df = pd.DataFrame(indices[:10].numpy(), columns=['index'])
words = pd.merge(df,
vocab[['index', 'word']],
how='left',
on='index')['word'].values.tolist()
logging.info(f"Topic {n}: {words}")
def train(self,
*,
raw_expr,
encoded_expr,
num_epochs=100,
batch_size=32,
learning_rate=1e-3,
eval_every=10,
test_proportion=0.05,
use_l1=False,
l1_lam=0):
seed = 2556
torch.manual_seed(seed)
pyro.set_rng_seed(seed)
pyro.clear_param_store()
logging.info('Validating data ...')
assert (raw_expr.shape == encoded_expr.shape)
read_depth = raw_expr.sum(-1)[:, np.newaxis]
encoded_expr = np.hstack([encoded_expr, np.log(read_depth)])
read_depth = torch.tensor(read_depth).to(self.device)
raw_expr = torch.tensor(raw_expr).to(self.device)
encoded_expr = torch.tensor(encoded_expr).to(self.device)
logging.info('Initializing model ...')
self.optimizer = Adam({"lr": 1e-3})
self.loss = TraceMeanField_ELBO()
if not use_l1:
logging.info('No L1 regularization.')
svi = SVI(self.model, self.guide, self.optimizer, loss=self.loss)
test_set = np.random.rand(read_depth.shape[0]) < test_proportion
train_set = ~test_set
logging.info("Training with {} cells, testing with {}.".format(
str(train_set.sum()), str(test_set.sum())))
logging.info('Training ...')
try:
for epoch in range(1, num_epochs + 1):
running_loss = 0.0
for batch in self.epoch_batch(raw_expr[train_set],
encoded_expr[train_set],
read_depth[train_set],
batch_size=batch_size):
if use_l1:
loss = self.custom_step(*batch)
else:
loss = svi.step(*batch)
running_loss += loss / batch_size
logging.info('Done epoch {}/{}. Training loss: {:.3e}'.format(
str(epoch), str(num_epochs), running_loss))
if (epoch % eval_every == 0
or epoch == num_epochs) and test_set.sum() > 0:
test_logp = self.evaluate(raw_expr[test_set],
encoded_expr[test_set],
read_depth[test_set])
logging.info('Test logp: {:.4e}'.format(test_logp))
except KeyboardInterrupt:
logging.error('Interrupted training.')
self.summarize_posterior(raw_expr, encoded_expr, read_depth)
return self
Evaluate DMEGP on PHYSIONET test dataset.
"""
# system configuration
torch.manual_seed(1)
# data configuration
dataloaders = fetch_dataloaders_PhysioNet(['val', 'test'], 'HR')
dataloader_val = dataloaders['val']
dataloader_test = dataloaders['test']
# load a model
model_path = './experiments/physionet/dmegp_physionet.pth'
cpt = torch.load(model_path)
embed_fn = MLP_embed(cpt['input_dim'], cpt['feature_dim'])
mean_fn = MLP(cpt['feature_dim'], cpt['hidden_dim'], cpt['output_dim'])
model = DMEGP(cpt['input_dim'],
feature_dim=cpt['feature_dim'],
mean_fn=mean_fn,
embed_fn=embed_fn)
model.load_state_dict(cpt['state_dict'])
elbo = TraceMeanField_ELBO()
loss_fn = elbo.differentiable_loss
# evaluate the model
rmse_val, loss_val = evaluate(model, dataloader_val, loss_fn,
cpt['n_adapt'], cpt['inner_lr'])
rmse_te, loss_te = evaluate(model, dataloader_test, loss_fn,
cpt['n_adapt'], cpt['inner_lr'])
print('[Val] RMSE={:05.3f} and loss={:05.3f}'.format(rmse_val, loss_val))
print('[Test] RMSE={:05.3f} and loss={:05.3f}'.format(rmse_te, loss_te))
def test_elbo_analytic_kl(self):
self.do_elbo_test(True, 3000, TraceMeanField_ELBO())
article_nav_bows = np.zeros(
(len(training_data.articles), len(training_data.nav_vocab)),
dtype=np.float32)
for article_id in range(len(training_data.article_navs)):
for nav_id in training_data.article_navs[article_id]:
article_nav_bows[article_id][nav_id] += 1
navLDA = NavEmbeddingLDA(device, nav_topics, float(args.dropout),
training_data).to(device)
optimizer = pyro.optim.Adam({"lr": float(args.lr)})
svi = SVI(navLDA.model,
navLDA.guide,
optimizer,
loss=TraceMeanField_ELBO())
elbos = []
num_batches = int(article_nav_bows.shape[0] / batch_size)
for epoch in range(int(args.num_epochs)):
article_ids = np.array((range(article_nav_bows.shape[0])))
np.random.shuffle(article_ids)
epoch_elbo = 0.0
for batch_id in range(num_batches):
batch_article_ids = article_ids[list(
range(batch_id * batch_size, (batch_id + 1) * batch_size))]
batch = torch.tensor(
article_nav_bows[batch_article_ids]).to(device)
elbos.append(
svi.step(batch, training_data.nav_embeddings,
def test_constant_kl():
model = lambda: pyro.sample("a", dist.Normal(0, 1.))
guide = tyxe.guides.AutoNormal(model)
elbo = TraceMeanField_ELBO()
assert elbo.loss(model, guide) == elbo.loss(model, guide)
def train(args, dataset):
"""
Train a model and guide to fit a dataset.
"""
counts = dataset["counts"]
num_stations = len(dataset["stations"])
train_size = args.truncate if args.truncate else len(counts)
logging.info(
"Training on {} stations over {}/{} hours, {} batches/epoch".format(
num_stations, train_size, len(counts),
int(math.ceil(train_size / args.batch_size))))
time_features = make_time_features(args, 0, len(counts))
control_features = (counts.max(1)[0] + counts.max(2)[0]).clamp(max=1)
logging.debug(
"On average {:0.1f}/{} stations are open at any one time".format(
control_features.sum(-1).mean(), num_stations))
features = torch.cat([time_features, control_features], -1)
feature_dim = features.size(-1)
logging.debug("feature_dim = {}".format(feature_dim))
metadata = {"args": args, "losses": [], "control": control_features}
torch.save(metadata, args.training_filename)
if args.device.startswith("cuda"):
torch.set_default_tensor_type('torch.cuda.FloatTensor')
def optim_config(module_name, param_name):
config = {
"lr": args.learning_rate,
"betas": (0.8, 0.99),
"lrd": 0.1**(1 / args.num_steps),
}
if param_name == "init_scale":
config["lr"] *= 0.1 # init_dist sees much less data per minibatch
return config
training_counts = counts[:args.truncate] if args.truncate else counts
data_size = len(training_counts)
model = Model(args, features, training_counts).to(device=args.device)
guide = Guide(args, features, training_counts).to(device=args.device)
optimizer = ClippedAdam(optim_config)
if args.funsor:
elbo = Funsor_ELBO(args)
elif args.analytic_kl:
elbo = TraceMeanField_ELBO()
else:
elbo = Trace_ELBO()
svi = SVI(model, guide, optimizer, elbo)
losses = []
forecaster = None
for step in range(args.num_steps):
begin_time = torch.randint(max(1, data_size - args.batch_size),
()).item()
end_time = min(data_size, begin_time + args.batch_size)
feature_batch = features[begin_time:end_time].to(device=args.device)
counts_batch = counts[begin_time:end_time].to(device=args.device)
loss = svi.step(feature_batch, counts_batch) / counts_batch.numel()
assert math.isfinite(loss), loss
losses.append(loss)
logging.debug("step {} loss = {:0.4g}".format(step, loss))
if step % 20 == 0:
# Save state every few steps.
pyro.get_param_store().save(args.param_store_filename)
metadata = {
"args": args,
"losses": losses,
"control": control_features
}
torch.save(metadata, args.training_filename)
forecaster = Forecaster(args, dataset, features, model, guide)
torch.save(forecaster, args.forecaster_filename)
if logging.Logger(None).isEnabledFor(logging.DEBUG):
init_scale = pyro.param("init_scale").data
trans_scale = pyro.param("trans_scale").data
trans_matrix = pyro.param("trans_matrix").data
eigs = trans_matrix.eig()[0].norm(dim=-1).sort(
descending=True).values
logging.debug("guide.diag_part = {}".format(
guide.diag_part.data.squeeze()))
logging.debug(
"init scale min/mean/max: {:0.3g} {:0.3g} {:0.3g}".format(
init_scale.min(), init_scale.mean(), init_scale.max()))
logging.debug(
"trans scale min/mean/max: {:0.3g} {:0.3g} {:0.3g}".format(
trans_scale.min(), trans_scale.mean(),
trans_scale.max()))
logging.debug("trans mat eig:\n{}".format(eigs))
return forecaster