def inference(train_x, train_y, num_epochs=10000):
svi = SVI(model,
guide,
optim.Adam({'lr': 0.001}),
loss=Trace_ELBO(),
num_samples=len(train_x))
for i in range(num_epochs):
elbo = svi.step(train_x, train_y)
if i % 1000 == 0:
print('Elbo loss : {}'.format(elbo))
print('pyro\'s Param Store')
for k, v in pyro.get_param_store().items():
print(k, v)
def update_noise_svi(self, observed_steady_state, initial_noise):
def guide(noise):
noise_terms = list(noise.keys())
mu_constraints = constraints.interval(-3., 3.)
sigma_constraints = constraints.interval(.0001, 3)
mu = {
k: pyro.param('{}_mu'.format(k),
tensor(0.),
constraint=mu_constraints)
for k in noise_terms
}
sigma = {
k: pyro.param('{}_sigma'.format(k),
tensor(1.),
constraint=sigma_constraints)
for k in noise_terms
}
for noise in noise_terms:
sample(noise, Normal(mu[noise], sigma[noise]))
observation_model = condition(self.noisy_model, observed_steady_state)
pyro.clear_param_store()
svi = SVI(model=observation_model,
guide=guide,
optim=SGD({
"lr": 0.001,
"momentum": 0.1
}),
loss=Trace_ELBO())
losses = []
num_steps = 1000
samples = defaultdict(list)
for t in range(num_steps):
losses.append(svi.step(initial_noise))
for noise in initial_noise.keys():
mu = '{}_mu'.format(noise)
sigma = '{}_sigma'.format(noise)
samples[mu].append(pyro.param(mu).item())
samples[sigma].append(pyro.param(sigma).item())
means = {k: statistics.mean(v) for k, v in samples.items()}
updated_noise = {
'N_Raf': Normal(means['N_Raf_mu'], means['N_Raf_sigma']),
'N_Mek': Normal(means['N_Mek_mu'], means['N_Mek_sigma']),
'N_Erk': Normal(means['N_Erk_mu'], means['N_Erk_sigma'])
}
return updated_noise, losses
def run_inference(self):
data = torch.tensor(parse_tdoc_file(self.observations_file_path,
self.documents_length,
self.words_number),
dtype=torch.float32).view(-1)
if self.create_ism_data_file:
self.ism.save_ism_data(data)
return 0
if self.use_ism:
seq = format_seq_file('./mutu_data/seq.txt')
self.initalized_motifs = self.ism.initialize_motifs(data, seq)
else:
self.initalized_motifs = torch.ones(self.latent_motifs_number, 1,
self.words_number,
self.relative_time_length)
pyro.clear_param_store()
adam_params = {"lr": 0.1, "betas": (0.9, 0.999)}
optimizer = pyro.optim.Adam(adam_params)
svi = SVI(self.model,
self.guide,
optimizer,
loss=Trace_ELBO(),
num_samples=self.n_samples)
for _ in tqdm(range(self.n_steps)):
svi.step(data)
motifs_starting_times_file_path = os.path.join(
self.work_dir, 'motifs_starting_times.npy')
motifs_file_path = os.path.join(self.work_dir, 'motifs.npy')
motifs_starting_times = pyro.param(
"q_motifs_starting_times").detach().numpy()
motifs = pyro.param("q_motifs").detach().numpy()
np.save(file=motifs_starting_times_file_path,
arr=motifs_starting_times)
np.save(file=motifs_file_path, arr=motifs)
self.dump_motifs_and_starting_times(motifs_starting_times, motifs)
if self.plot_results:
self.plot_motifs_and_starting_times(motifs)
def fit(self, data_loaders, num_epochs=10):
train_loader, test_loader = data_loaders
batch_size = train_loader.batch_size
optimizer = Adam({'lr': 0.0009})
elbo = Trace_ELBO()
svi = SVI(self.model, self.guide, optimizer, elbo)
for epoch in range(num_epochs):
loss = 0.
iterations = 0
for x, y in tqdm(train_loader):
if x.size() == (batch_size, 1, 28, 28):
loss += svi.step(x.reshape(batch_size, -1), y)
iterations += 1
print(epoch, ' Epoch loss : ', loss / (iterations * batch_size))
def inference_svi(self, data, steps=3000, lr=0.01):
self.inference_method = "svi"
pyro.clear_param_store()
self.optimizer = Adam({"lr": lr, "betas": (0.90, 0.999)})
self.svi = SVI(self.model,
self.multi_norm_guide(),
self.optimizer,
loss=Trace_ELBO())
self.history = {"losses": []}
data = torch.tensor(data).float()
bar = trange(steps)
for i in bar:
loss = self.svi.step(data)
if (i + 1) % 100 == 1:
bar.write("Now step %d completed, loss is %.4f" % (i, loss))
self.history["losses"].append(loss)
def __init__(self,model, guide=None, epochs=500, lr=0.05, cuda=False):
self.model=model
self.cuda=cuda
if guide != None:
self.guide=guide
else:
self.guide=AutoDiagonalNormal(model)
self.optimizer=optim.Adam({"lr": lr})
self.svi=SVI(self.model,
self.guide,
self.optimizer,
loss=Trace_ELBO())
self.epochs=epochs
def test_pyrocov_smoke(model, Guide, backend):
T, P, S, F = 3, 4, 5, 6
dataset = {
"features": torch.randn(S, F),
"local_time": torch.randn(T, P),
"weekly_strains": torch.randn(T, P, S).exp().round(),
}
guide = Guide(model, backend=backend)
svi = SVI(model, guide, ClippedAdam({"lr": 1e-8}), Trace_ELBO())
for step in range(2):
with xfail_if_not_implemented():
svi.step(dataset)
guide(dataset)
predictive = Predictive(model, guide=guide, num_samples=2)
predictive(dataset)
def test_posterior_predictive_svi_one_hot():
pseudocounts = torch.ones(3) * 0.1
true_probs = torch.tensor([0.15, 0.6, 0.25])
classes = dist.OneHotCategorical(true_probs).sample((10000, ))
guide = AutoDelta(one_hot_model)
svi = SVI(one_hot_model, guide, optim.Adam(dict(lr=0.1)), Trace_ELBO())
for i in range(1000):
svi.step(pseudocounts, classes=classes)
posterior_samples = Predictive(guide,
num_samples=10000).get_samples(pseudocounts)
posterior_predictive = Predictive(one_hot_model, posterior_samples)
marginal_return_vals = posterior_predictive.get_samples(
pseudocounts)["obs"]
assert_close(marginal_return_vals.mean(dim=0),
true_probs.unsqueeze(0),
rtol=0.1)
def test_laplace_linear_model(linear_model, one_point_design):
pyro.set_rng_seed(42)
pyro.clear_param_store()
# You can use 1 final sample here because linear models have a posterior entropy that is independent of `y`
estimated_eig = laplace_eig(linear_model,
one_point_design,
"y",
"w",
guide=laplace_guide,
num_steps=250,
final_num_samples=1,
optim=optim.Adam({"lr": 0.05}),
loss=Trace_ELBO().differentiable_loss)
expected_eig = linear_model_ground_truth(linear_model, one_point_design,
"y", "w")
assert_equal(estimated_eig, expected_eig, prec=5e-2)
def test_subsample_guide(auto_class, init_fn):
# The model from tutorial/source/easyguide.ipynb
def model(batch, subsample, full_size):
num_time_steps = len(batch)
result = [None] * num_time_steps
drift = pyro.sample("drift", dist.LogNormal(-1, 0.5))
plate = pyro.plate("data", full_size, subsample=subsample)
assert plate.size == 50
with plate:
z = 0.
for t in range(num_time_steps):
z = pyro.sample("state_{}".format(t), dist.Normal(z, drift))
result[t] = pyro.sample("obs_{}".format(t),
dist.Bernoulli(logits=z),
obs=batch[t])
return torch.stack(result)
def create_plates(batch, subsample, full_size):
return pyro.plate("data", full_size, subsample=subsample)
if auto_class == AutoGuideList:
guide = AutoGuideList(model, create_plates=create_plates)
guide.add(AutoDelta(poutine.block(model, expose=["drift"])))
guide.add(AutoNormal(poutine.block(model, hide=["drift"])))
else:
guide = auto_class(model, create_plates=create_plates)
full_size = 50
batch_size = 20
num_time_steps = 8
pyro.set_rng_seed(123456789)
data = model([None] * num_time_steps, torch.arange(full_size), full_size)
assert data.shape == (num_time_steps, full_size)
pyro.get_param_store().clear()
pyro.set_rng_seed(123456789)
svi = SVI(model, guide, Adam({"lr": 0.02}), Trace_ELBO())
for epoch in range(2):
beg = 0
while beg < full_size:
end = min(full_size, beg + batch_size)
subsample = torch.arange(beg, end)
batch = data[:, beg:end]
beg = end
svi.step(batch, subsample, full_size=full_size)
def train(
self,
max_epochs: int = 30000,
batch_size: int = None,
train_size: float = 1,
lr: float = 0.002,
num_particles: int = 1,
scale_elbo: float = 1.0,
**kwargs,
):
"""Train the model with useful defaults
Parameters
----------
max_epochs
Number of passes through the dataset. If `None`, defaults to
`np.min([round((20000 / n_cells) * 400), 400])`
train_size
Size of training set in the range [0.0, 1.0]. Use all data points in training because
we need to estimate cell abundance at all locations.
batch_size
Minibatch size to use during training. If `None`, no minibatching occurs and all
data is copied to device (e.g., GPU).
lr
Optimiser learning rate (default optimiser is :class:`~pyro.optim.ClippedAdam`).
Specifying optimiser via plan_kwargs overrides this choice of lr.
kwargs
Other arguments to scvi.model.base.PyroSviTrainMixin().train() method
"""
kwargs["max_epochs"] = max_epochs
kwargs["batch_size"] = batch_size
kwargs["train_size"] = train_size
kwargs["lr"] = lr
if "plan_kwargs" not in kwargs.keys():
kwargs["plan_kwargs"] = dict()
if getattr(self.module.model, "discrete_variables", None) and (len(self.module.model.discrete_variables) > 0):
kwargs["plan_kwargs"]["loss_fn"] = TraceEnum_ELBO(num_particles=num_particles)
else:
kwargs["plan_kwargs"]["loss_fn"] = Trace_ELBO(num_particles=num_particles)
if scale_elbo != 1.0:
if scale_elbo is None:
scale_elbo = 1.0 / (self.summary_stats["n_cells"] * self.summary_stats["n_genes"])
kwargs["plan_kwargs"]["scale_elbo"] = scale_elbo
super().train(**kwargs)
def inference(train_x, train_y, num_epochs=2000):
svi = SVI(model,
guide,
optim.Adam({'lr': 0.005}),
loss=Trace_ELBO(),
num_samples=1000)
for i in range(num_epochs):
elbo = svi.step(train_x, train_y)
if i % 200 == 0:
print('Elbo loss : {}'.format(elbo))
svi_posterior = svi.run(train_x, train_y)
sites = ['w', 'b', 'sigma']
for site, values in summary(svi_posterior, sites).items():
print("Site: {}".format(site))
print(values, "\n")
def main(model, guide, args):
# init
if args.seed is not None: pyro.set_rng_seed(args.seed)
logger = get_logger(args.log, __name__)
logger.info(args)
# load data
train_loader, test_loader = setup_data_loaders(batch_size=256)
# setup svi
pyro.clear_param_store()
# # WL: edited to make SCORE produce no NaNs. =====
# # lr value:
# # - original value in vae: 1.0e-3 --- SCORE produces NaNs at some point.
# # - default of Adam(..) : 1.0e-3
# # - current value : 1.0e-4 --- SCORE produces no NaNs.
# learning_rate = 1.0e-4
# # ===============================================
opt = optim.Adam({"lr": args.learning_rate})
elbo = Trace_ELBO()
svi = SVI(model.main, guide.main, opt, loss=elbo)
# # train (init)
# loss_avg = evaluate(svi, train_loader)
# param_state = copy.deepcopy(pyro.get_param_store().get_state())
# elbo_l = [-loss_avg]
# param_state_l = [param_state]
# train
times = [time.time()]
logger.info(f"\nepoch\t" + "elbo\t" + "time(sec)")
for i in range(1, args.num_epochs + 1):
loss_avg = train_epoch(svi, train_loader)
# elbo_l.append(-loss_avg)
#
# if (i+1) % param_freq == 0:
# param_state = copy.deepcopy(pyro.get_param_store().get_state())
# param_state_l.append(param_state)
if (args.eval_frequency > 0
and i % args.eval_frequency == 0) or (i == 1):
times.append(time.time())
logger.info(f"{i:06d}\t"
f"{-loss_avg:.4f}\t"
f"{times[-1]-times[-2]:.3f}")
def train(self, train_data, len_dataset, num_iterations):
# Treina o algoritmo, usando inferência variacional com o
# otimizador Adam, ambos implementados no Pyro
optim = Adam({"lr": 0.01})
svi = SVI(self.model, self.guide, optim, loss=Trace_ELBO())
loss = 0
losses = []
for j in range(num_iterations):
loss = 0
for data in train_data:
# calculate the loss and take a gradient step
loss += svi.step(data[0].float(), data[1])
total_epoch_loss_train = loss / len_dataset
losses.append(total_epoch_loss_train)
#print("Epoca: ", j, " Perda: ", total_epoch_loss_train)
return losses
def test_posterior_predictive_svi_auto_delta_guide(parallel):
true_probs = torch.ones(5) * 0.7
num_trials = torch.ones(5) * 1000
num_success = dist.Binomial(num_trials, true_probs).sample()
conditioned_model = poutine.condition(model, data={"obs": num_success})
guide = AutoDelta(conditioned_model)
svi = SVI(conditioned_model, guide, optim.Adam(dict(lr=1.0)), Trace_ELBO())
for i in range(1000):
svi.step(num_trials)
posterior_predictive = Predictive(model,
guide=guide,
num_samples=10000,
parallel=parallel)
marginal_return_vals = posterior_predictive.get_samples(num_trials)["obs"]
assert_close(marginal_return_vals.mean(dim=0),
torch.ones(5) * 700,
rtol=0.05)
def get_pyro_model(return_all=True):
nn_model = NN_Model(input_size=3, hidden_size=2, output_size=1)
model = model_fn(nn_model)
guide = guide_fn(nn_model)
AdamArgs = {'lr': 3e-3}
optimizer = torch.optim.Adam
scheduler = pyro.optim.ExponentialLR({
'optimizer': optimizer,
'optim_args': AdamArgs,
'gamma': 0.99995
})
svi = SVI(model, guide, scheduler, loss=Trace_ELBO(), num_samples=1000)
if return_all:
return svi, model, guide
else:
return svi
def run_svi(model, guide, x, y, optimizer=pyro.optim.Adam({"lr": 0.03}),
loss=Trace_ELBO(), num_iters=5000, verbose=True):
pyro.clear_param_store()
svi = SVI(model,
guide,
optimizer,
loss)
elbos = []
params = []
for i in range(num_iters):
elbo = svi.step(x, y)
if verbose and i % 500 == 0:
logging.info("Elbo loss: {}".format(elbo))
elbos.append(elbo)
params.append(dict(copy(pyro.get_param_store())))
return elbos, params
def fit(self, optim=Adam({'lr': 1e-3}), loss=Trace_ELBO(num_particles=1), max_iter=5000, random_instance=None):
svi = SVI(self.model, self.guide, optim=optim, loss=loss)
with trange(max_iter) as t:
for i in t:
t.set_description(f'迭代：{i}')
svi.step(self.data)
loss = svi.evaluate_loss(self.data)
with torch.no_grad():
postfix_kwargs = {}
if random_instance is not None:
g = pyro.param('g')
s = pyro.param('s')
postfix_kwargs.update({
'g': '{0}'.format((g - random_instance.g).abs().mean()),
's': '{0}'.format((s - random_instance.s).abs().mean())
})
t.set_postfix(loss=loss, **postfix_kwargs)
def test_svi_smoke():
class Model(PyroModule):
def __init__(self):
super().__init__()
self.loc = nn.Parameter(torch.zeros(2))
self.scale = PyroParam(torch.ones(2),
constraint=constraints.positive)
self.z = PyroSample(
lambda self: dist.Normal(self.loc, self.scale).to_event(1))
def forward(self, data):
loc, log_scale = self.z.unbind(-1)
with pyro.plate("data"):
pyro.sample("obs", dist.Cauchy(loc, log_scale.exp()), obs=data)
class Guide(PyroModule):
def __init__(self):
super().__init__()
self.loc = nn.Parameter(torch.zeros(2))
self.scale = PyroParam(torch.ones(2),
constraint=constraints.positive)
self.z = PyroSample(
lambda self: dist.Normal(self.loc, self.scale).to_event(1))
def forward(self, *args, **kwargs):
return self.z
data = torch.randn(5)
model = Model()
trace = poutine.trace(model).get_trace(data)
assert "loc" in trace.nodes.keys()
assert trace.nodes["loc"]["type"] == "param"
assert "scale" in trace.nodes
assert trace.nodes["scale"]["type"] == "param"
guide = Guide()
trace = poutine.trace(guide).get_trace(data)
assert "loc" in trace.nodes.keys()
assert trace.nodes["loc"]["type"] == "param"
assert "scale" in trace.nodes
assert trace.nodes["scale"]["type"] == "param"
optim = Adam({"lr": 0.01})
svi = SVI(model, guide, optim, Trace_ELBO())
for step in range(3):
svi.step(data)
def do_test_fixedness(self, fixed_parts):
pyro.clear_param_store()
def model():
alpha_p_log = pyro.param("alpha_p_log", self.alpha_p_log_0.clone())
beta_p_log = pyro.param("beta_p_log", self.beta_p_log_0.clone())
alpha_p, beta_p = torch.exp(alpha_p_log), torch.exp(beta_p_log)
lambda_latent = pyro.sample("lambda_latent",
dist.Gamma(alpha_p, beta_p))
pyro.sample("obs", dist.Poisson(lambda_latent), obs=self.data)
return lambda_latent
def guide():
alpha_q_log = pyro.param("alpha_q_log", self.alpha_q_log_0.clone())
beta_q_log = pyro.param("beta_q_log", self.beta_q_log_0.clone())
alpha_q, beta_q = torch.exp(alpha_q_log), torch.exp(beta_q_log)
pyro.sample("lambda_latent", dist.Gamma(alpha_q, beta_q))
def per_param_args(param_name):
if "model" in fixed_parts and "p_" in param_name:
return {"lr": 0.0}
if "guide" in fixed_parts and "q_" in param_name:
return {"lr": 0.0}
return {"lr": 0.01}
adam = optim.Adam(per_param_args)
svi = SVI(model, guide, adam, loss=Trace_ELBO())
for _ in range(3):
svi.step()
model_unchanged = (torch.equal(
pyro.param("alpha_p_log").data,
self.alpha_p_log_0)) and (torch.equal(
pyro.param("beta_p_log").data, self.beta_p_log_0))
guide_unchanged = (torch.equal(
pyro.param("alpha_q_log").data,
self.alpha_q_log_0)) and (torch.equal(
pyro.param("beta_q_log").data, self.beta_q_log_0))
model_changed = not model_unchanged
guide_changed = not guide_unchanged
error = ("model" in fixed_parts
and model_changed) or ("guide" in fixed_parts
and guide_changed)
return not error
def fit(
self,
df,
max_iter=6000,
patience=200,
optimiser_settings={"lr": 1.0e-2},
elbo_kwargs={"num_particles": 5},
):
teams = sorted(list(set(df["home_team"]) | set(df["away_team"])))
home_team = df["home_team"].values
away_team = df["away_team"].values
home_goals = torch.tensor(df["home_goals"].values, dtype=torch.float32)
away_goals = torch.tensor(df["away_goals"].values, dtype=torch.float32)
gameweek = ((df["date"] - df["date"].min()).dt.days // 7).values
self.team_to_index = {team: i for i, team in enumerate(teams)}
self.index_to_team = {
value: key
for key, value in self.team_to_index.items()
}
self.n_teams = len(teams)
self.min_date = df["date"].min()
conditioned_model = condition(self.model,
data={
"home_goals": home_goals,
"away_goals": away_goals
})
guide = AutoDiagonalNormal(conditioned_model)
optimizer = Adam(optimiser_settings)
elbo = Trace_ELBO(**elbo_kwargs)
svi = SVI(conditioned_model, guide, optimizer, loss=elbo)
pyro.clear_param_store()
fitted_svi, losses = early_stopping(svi,
home_team,
away_team,
gameweek,
max_iter=max_iter,
patience=patience)
self.guide = guide
return losses
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
def train_model(pmf):
def model(data):
# sample f from the prior
# Probabilities are generated by the pmf
f = pyro.sample("latent_fairness", pmf)
f2 = dist.Bernoulli(f)
for i in range(len(data)):
s = pyro.sample("obs_{}".format(i), f2, obs=data[i])
def guide(data):
alpha_q = pyro.param("alpha_q",
torch.tensor(15.0),
constraint=constraints.positive)
beta_q = pyro.param("beta_q",
torch.tensor(15.0),
constraint=constraints.positive)
# sample latent_fairness from the distribution Beta(alpha_q, beta_q)
pyro.sample("latent_fairness", dist.Beta(alpha_q, beta_q))
adam_params = {"lr": 0.0005, "betas": (0.90, 0.999)}
optimizer = ClippedAdam(adam_params)
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
for step in range(n_steps):
loss = svi.step(data)
if step % 100 == 0:
logging.info(".")
logging.info("Elbo loss: {}".format(loss))
# grab the learned variational parameters
a_q = pyro.param("alpha_q").item()
b_q = pyro.param("beta_q").item()
inferred_mean = a_q / (a_q + b_q)
# compute inferred standard deviation
factor = b_q / (a_q * (1.0 + a_q + b_q))
inferred_std = inferred_mean * math.sqrt(factor)
print("\nbased on the data and our prior belief, the fairness " +
"of the coin is %.3f +- %.3f" % (inferred_mean, inferred_std))
beta_posterior = torch.distributions.beta.Beta(a_q, b_q)
posterior = torch.distributions.bernoulli.Bernoulli(
beta_posterior.sample())
logging.info("Sampling:{}".format(posterior.sample()))
def main(args):
pyro.set_rng_seed(1219)
pyro.clear_param_store()
data = pickle.load(open('data/results.pkl', 'rb'))
data = torch.tensor(data)
train_set, test_set = data[:30000], data[30000:]
train_set = torch.utils.data.TensorDataset(train_set)
test_set = torch.utils.data.TensorDataset(test_set)
kwargs = {'num_workers': 1, 'pin_memory': args.use_cuda}
train_loader = torch.utils.data.DataLoader(dataset=train_set,
batch_size=args.batchsize,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(dataset=test_set,
batch_size=args.batchsize,
shuffle=False,
**kwargs)
br = BayesianRanking(2000)
optimizer = Adam({"lr": args.lr})
svi = SVI(br.model, br.guide, optimizer, loss=Trace_ELBO())
train_elbo = []
test_elbo = []
for epoch in range(args.epochs):
total_epoch_loss_train = train(svi,
train_loader,
use_cuda=args.use_cuda)
train_elbo.append(-total_epoch_loss_train)
print("[epoch %03d] average training loss: %.4f" %
(epoch, total_epoch_loss_train))
if epoch % args.test_freq == 0:
# report test diagnostics
total_epoch_loss_test = evaluate(svi,
test_loader,
use_cuda=args.use_cuda)
test_elbo.append(-total_epoch_loss_test)
print("[epoch %03d] average test loss: %.4f" %
(epoch, total_epoch_loss_test))
mu = pickle.load(open('./data/mu_gt.pkl', 'rb'))
pred_mu_q = pyro.param('mu_q').detach().numpy().squeeze()
print(ss.spearmanr(mu, pred_mu_q))
def load(self, path):
"""Load model from path
Args:
load_path (string): Path to saved model
"""
model_path = path + '_model'
opt_path = path + '_opt'
guide_path = path + '_guide'
self.net = torch.load(model_path)
pyro.get_param_store().load(path + '_params')
self.optim = Adam({"lr": self.learning_rate})
self.optim.load(opt_path)
self.guide = AutoDiagonalNormal(self.model)
self.guide = torch.load(guide_path)
self.svi = SVI(self.net, self.guide, self.optim, loss=Trace_ELBO())
def test_posterior_predictive_svi_manual_guide(parallel):
true_probs = torch.ones(5) * 0.7
num_trials = torch.ones(5) * 1000
num_success = dist.Binomial(num_trials, true_probs).sample()
conditioned_model = poutine.condition(model, data={"obs": num_success})
elbo = Trace_ELBO(num_particles=100, vectorize_particles=True)
svi = SVI(conditioned_model, beta_guide, optim.Adam(dict(lr=1.0)), elbo)
for i in range(1000):
svi.step(num_trials)
posterior_predictive = Predictive(
model,
guide=beta_guide,
num_samples=10000,
parallel=parallel,
return_sites=["_RETURN"],
)
marginal_return_vals = posterior_predictive(num_trials)["_RETURN"]
assert_close(marginal_return_vals.mean(dim=0), torch.ones(5) * 700, rtol=0.05)
def train(model):
pyro.clear_param_store()
model(x)
guide = AutoDiagonalNormal(model)
guide(x)
adam = optim.Adam({"lr": 0.1})
svi = SVI(model, guide, adam, loss=Trace_ELBO())
for _ in range(1000):
svi.step(x)
import matplotlib.pyplot as plt
med = guide.median()['trans_score.weight'].detach()
plt.imshow(med)
return med
def test_posterior_predictive_svi_auto_diag_normal_guide():
true_probs = torch.ones(5) * 0.7
num_trials = torch.ones(5) * 1000
num_success = dist.Binomial(num_trials, true_probs).sample()
conditioned_model = poutine.condition(model, data={"obs": num_success})
opt = optim.Adam(dict(lr=0.1))
loss = Trace_ELBO()
guide = AutoDiagonalNormal(conditioned_model)
svi_run = SVI(conditioned_model,
guide,
opt,
loss,
num_steps=1000,
num_samples=100).run(num_trials)
posterior_predictive = TracePredictive(model, svi_run,
num_samples=10000).run(num_trials)
marginal_return_vals = posterior_predictive.marginal().empirical["_RETURN"]
assert_close(marginal_return_vals.mean, torch.ones(5) * 700, rtol=0.05)
def test_posterior_predictive_svi_one_hot():
pseudocounts = torch.ones(3) * 0.1
true_probs = torch.tensor([0.15, 0.6, 0.25])
classes = dist.OneHotCategorical(true_probs).sample((10000, ))
opt = optim.Adam(dict(lr=0.1))
loss = Trace_ELBO()
guide = AutoDelta(one_hot_model)
svi_run = SVI(one_hot_model,
guide,
opt,
loss,
num_steps=1000,
num_samples=1000).run(pseudocounts, classes=classes)
posterior_predictive = TracePredictive(one_hot_model,
svi_run,
num_samples=10000).run(pseudocounts)
marginal_return_vals = posterior_predictive.marginal().empirical["_RETURN"]
assert_close(marginal_return_vals.mean, true_probs.unsqueeze(0), rtol=0.1)
def _train_svi(self, train_loader, epochs, lr, savedir, device):
print("\n == fullBNN SVI training ==")
optimizer = pyro.optim.Adam({"lr":lr})
elbo = Trace_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)
loss += svi.step(x_data=x_batch, y_data=y_batch.argmax(dim=-1))
outputs = self.forward(x_batch, training=True, avg_posterior=False).to(device)
predictions = outputs.argmax(-1)
labels = y_batch.argmax(-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(savedir)
plot_loss_accuracy(dict={'loss':loss_list, 'accuracy':accuracy_list},
path=os.path.join(savedir, self.name+"_training.png"))
