def main(args):
pyro.enable_validation(__debug__)
if args.cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
# Generate synthetic data.
pyro.set_rng_seed(args.seed)
x_train, t_train, y_train, _ = generate_data(args)
# Train.
pyro.set_rng_seed(args.seed)
pyro.clear_param_store()
cevae = CEVAE(feature_dim=args.feature_dim,
latent_dim=args.latent_dim,
hidden_dim=args.hidden_dim,
num_layers=args.num_layers,
num_samples=10)
cevae.fit(x_train, t_train, y_train,
num_epochs=args.num_epochs,
batch_size=args.batch_size,
learning_rate=args.learning_rate,
learning_rate_decay=args.learning_rate_decay,
weight_decay=args.weight_decay)
# Evaluate.
x_test, t_test, y_test, true_ite = generate_data(args)
true_ate = true_ite.mean()
print("true ATE = {:0.3g}".format(true_ate.item()))
naive_ate = y_test[t_test == 1].mean() - y_test[t_test == 0].mean()
print("naive ATE = {:0.3g}".format(naive_ate))
if args.jit:
cevae = cevae.to_script_module()
est_ite = cevae.ite(x_test)
est_ate = est_ite.mean()
print("estimated ATE = {:0.3g}".format(est_ate.item()))
def main(args):
assert pyro.__version__ >= "0.4.1"
pyro.enable_validation(__debug__)
pyro.set_rng_seed(args.seed)
dataset = load_hourly_od(args)
if args.tiny:
dataset["stations"] = dataset["stations"][:args.tiny]
dataset["counts"] = dataset["counts"][:, :args.tiny, :args.tiny]
forecaster = train(args, dataset)
if forecaster is None:
return
window_begin = max(0, args.truncate - args.batch_size)
window_end = args.truncate
truth = dataset['counts'][window_end:window_end + args.forecast_hours]
forecast = forecaster(window_begin,
window_end,
args.forecast_hours,
num_samples=args.num_samples)
assert forecast.shape == (args.num_samples, ) + truth.shape
log_prob = forecaster.log_prob(window_begin, window_end,
truth) / truth.numel()
torch.save({
'forecast': forecast,
'truth': truth,
'log_prob': log_prob,
}, args.forecast_filename)
def main(args):
pyro.set_rng_seed(0)
pyro.enable_validation(__debug__)
optim = Adam({"lr": 0.1})
inference = SVI(model, guide, optim, loss=Trace_ELBO())
# Data is an arbitrary json-like structure with tensors at leaves.
one = torch.tensor(1.0)
data = {
"foo": one,
"bar": [0 * one, 1 * one, 2 * one],
"baz": {
"noun": {
"concrete": 4 * one,
"abstract": 6 * one,
},
"verb": 2 * one,
},
}
print('Step\tLoss')
loss = 0.0
for step in range(args.num_epochs):
loss += inference.step(data)
if step and step % 10 == 0:
print('{}\t{:0.5g}'.format(step, loss))
loss = 0.0
print('Parameters:')
for name, value in sorted(pyro.get_param_store().items()):
print('{} = {}'.format(name, value.detach().cpu().numpy()))
def main():
inputs,calendar = load_input()
logger.info('Inference')
covariates, covariate_dim, data = inputs.values()
data, covariates = map(jax_to_torch,[data,covariates])
data = torch.log(1+data.double())
assert pyro.__version__.startswith('1.3.1')
pyro.enable_validation(True)
T0 = 0 # begining
T2 = data.size(-2) # end
T1 = T2 - 500 # train/test split
pyro.set_rng_seed(1)
pyro.clear_param_store()
data = data.permute(-2,-1)
covariates = covariates.reshape(data.size(-1),T2,-1)
# covariates = torch.zeros(len(data), 0) # empty
forecaster = Forecaster(Model4(), data[:T1], covariates[:,:T1], learning_rate=0.09,num_steps=2000)
samples = forecaster(data[:T1], covariates[:,:T2], num_samples=336)
samples.clamp_(min=0) # apply domain knowledge: the samples must be positive
p10, p50, p90 = quantile(samples[:, 0], [0.1, 0.5, 0.9]).squeeze(-1)
crps = eval_crps(samples, data[T1:T2])
print(samples.shape, p10.shape)
fig, axes = plt.subplots(data.size(-1), 1, figsize=(9, 10), sharex=True)
plt.subplots_adjust(hspace=0)
axes[0].set_title("Sales (CRPS = {:0.3g})".format(crps))
for i, ax in enumerate(axes):
ax.fill_between(torch.arange(T1, T2), p10[:, i], p90[:, i], color="red", alpha=0.3)
ax.plot(torch.arange(T1, T2), p50[:, i], 'r-', lw=1, label='forecast')
ax.plot(torch.arange(T0, T2),data[: T2, i], 'k-', lw=1, label='truth')
ax.set_ylabel(f"item: {i}")
axes[0].legend(loc="best")
plt.show()
plt.savefig('figures/pyro_forecast.png')
def main(args):
pyro.set_rng_seed(args.seed)
pyro.enable_validation(__debug__)
model = SimpleHarmonicModel(args.process_noise, args.measurement_noise)
guide = SimpleHarmonicModel_Guide(model)
smc = SMCFilter(model,
guide,
num_particles=args.num_particles,
max_plate_nesting=0)
logging.info('Generating data')
zs, ys = generate_data(args)
logging.info('Filtering')
smc.init(initial=torch.tensor([1., 0.]))
for y in ys[1:]:
smc.step(y)
logging.info('Marginals')
empirical = smc.get_empirical()
for t in range(1, 1 + args.num_timesteps):
z = empirical["z_{}".format(t)]
logging.info("{}\t{}\t{}\t{}".format(t, zs[t], z.mean, z.variance))
def __init__(self, hparams, pyro_model:BaseSEM):
super().__init__()
self.pyro_model = pyro_model
hparams.experiment = self.__class__.__name__
hparams.model = pyro_model.__class__.__name__
self.hparams = hparams
self.train_batch_size = hparams.train_batch_size
self.test_batch_size = hparams.test_batch_size
if hasattr(hparams, 'num_sample_particles'):
self.pyro_model._gen_counterfactual = partial(self.pyro_model.counterfactual, num_particles=self.hparams.num_sample_particles)
else:
self.pyro_model._gen_counterfactual = self.pyro_model.counterfactual
if hparams.validate:
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.autograd.set_detect_anomaly(self.hparams.validate)
pyro.enable_validation()
resize = None if self.hparams.resize == (0,0) else self.hparams.resize
train_crop_type = self.hparams.train_crop_type if hasattr(self.hparams, 'train_crop_type') else 'random'
crop_size = self.hparams.crop_size if hasattr(self.hparams, 'crop_size') else (224, 224)
self.calabresi_train = CalabresiDataset(self.hparams.train_csv, crop_size=crop_size, crop_type=train_crop_type, resize=resize) # noqa: E501
self.calabresi_val = CalabresiDataset(self.hparams.valid_csv, crop_size=crop_size, crop_type='center', resize=resize)
self.calabresi_test = CalabresiDataset(self.hparams.test_csv, crop_size=crop_size, crop_type='center', resize=resize)
def main(args):
# Init Pyro
pyro.enable_validation(True)
pyro.clear_param_store()
# Load meta-data for all models and select model based on command arguments
models = pyro_models.load()
#model_dict = select_model(args, models)
model_dict = models['arm.earnings_latin_square']
#model_dict = models['arm.election88_ch14']
# Define model/data/guide
model = model_dict['model']
data = pyro_models.data(model_dict)
guide = AutoDelta(model)
# Perform variational inference
ess = ESS(vectorize_particles=True, num_inner=1000, num_outer=10)
svi = SVI(model, guide, optim.Adam({'lr': 0.1}), loss=Trace_ELBO())
for i in range(args.num_epochs):
params = {}
ess_val = ess.loss(model, guide, data, {})
loss = svi.step(data, params)
print('loss', loss, 'ess', ess_val)
sys.exit()
def main(args):
logging.info('Generating data')
pyro.set_rng_seed(0)
pyro.clear_param_store()
pyro.enable_validation(__debug__)
# We can generate synthetic data directly by calling the model.
true_topic_weights, true_topic_words, data = model(args=args)
# We'll train using SVI.
logging.info('-' * 40)
logging.info('Training on {} documents'.format(args.num_docs))
predictor = make_predictor(args)
guide = functools.partial(parametrized_guide, predictor)
Elbo = JitTraceEnum_ELBO if args.jit else TraceEnum_ELBO
elbo = Elbo(max_plate_nesting=2)
optim = ClippedAdam({'lr': args.learning_rate})
svi = SVI(model, guide, optim, elbo)
logging.info('Step\tLoss')
for step in range(args.num_steps):
loss = svi.step(data, args=args, batch_size=args.batch_size)
if step % 10 == 0:
logging.info('{: >5d}\t{}'.format(step, loss))
loss = elbo.loss(model, guide, data, args=args)
logging.info('final loss = {}'.format(loss))
def main(args):
pyro.set_rng_seed(args.seed)
pyro.enable_validation(__debug__)
model = SimpleHarmonicModel(args.process_noise, args.measurement_noise)
guide = SimpleHarmonicModel_Guide(model)
smc = SMCFilter(model,
guide,
num_particles=args.num_particles,
max_plate_nesting=0)
logging.info("Generating data")
zs, ys = generate_data(args)
logging.info("Filtering")
smc.init(initial=torch.tensor([1., 0.]))
for y in ys[1:]:
smc.step(y)
logging.info("At final time step:")
z = smc.get_empirical()["z"]
logging.info("truth: {}".format(zs[-1]))
logging.info("mean: {}".format(z.mean))
logging.info("std: {}".format(z.variance**0.5))
def main(args):
pyro.set_rng_seed(0)
pyro.enable_validation()
optim = Adam({"lr": 0.1})
inference = SVI(model, guide, optim, loss=Trace_ELBO())
# Data is an arbitrary json-like structure with tensors at leaves.
one = torch.tensor(1.0)
data = {
"foo": one,
"bar": [0 * one, 1 * one, 2 * one],
"baz": {
"noun": {
"concrete": 4 * one,
"abstract": 6 * one,
},
"verb": 2 * one,
},
}
print('Step\tLoss')
loss = 0.0
for step in range(args.num_epochs):
loss += inference.step(data)
if step and step % 10 == 0:
print('{}\t{:0.5g}'.format(step, loss))
loss = 0.0
print('Parameters:')
for name in sorted(pyro.get_param_store().get_all_param_names()):
print('{} = {}'.format(name, pyro.param(name).detach().cpu().numpy()))
def __init__(self, hparams, pyro_model: BaseSEM):
super().__init__()
self.pyro_model = pyro_model
hparams.experiment = self.__class__.__name__
hparams.model = pyro_model.__class__.__name__
self.hparams = hparams
self.data_dir = hparams.data_dir
self.train_batch_size = hparams.train_batch_size
self.test_batch_size = hparams.test_batch_size
if hasattr(hparams, 'num_sample_particles'):
self.pyro_model._gen_counterfactual = partial(
self.pyro_model.counterfactual,
num_particles=self.hparams.num_sample_particles)
else:
self.pyro_model._gen_counterfactual = self.pyro_model.counterfactual
if hparams.validate:
import random
torch.manual_seed(0)
np.random.seed(0)
random.seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.autograd.set_detect_anomaly(self.hparams.validate)
pyro.enable_validation()
def main(args):
logging.info(f"CUDA enabled: {torch.cuda.is_available()}")
logging.info('Generating data')
pyro.set_rng_seed(0)
pyro.clear_param_store()
pyro.enable_validation(__debug__)
# We can generate synthetic data directly by calling the model.
data = model(args=args)
doc_word_data = data["doc_word_data"]
category_data = data["category_data"]
# We'll train using SVI.
logging.info('-' * 40)
logging.info('Training on {} documents'.format(args.num_docs))
Elbo = JitTraceEnum_ELBO if args.jit else TraceEnum_ELBO
elbo = Elbo(max_plate_nesting=2)
optim = ClippedAdam({'lr': args.learning_rate})
svi = SVI(model, parametrized_guide, optim, elbo)
logging.info('Step\tLoss')
for step in tqdm(range(args.num_steps)):
loss = svi.step(doc_word_data=doc_word_data,
category_data=category_data,
args=args,
batch_size=args.batch_size)
if step % 10 == 0:
logging.info('{: >5d}\t{}'.format(step, loss))
loss = elbo.loss(model,
parametrized_guide,
doc_word_data=doc_word_data,
category_data=category_data,
args=args)
logging.info('final loss = {}'.format(loss))
print("debug string")
def main(args):
funsor.set_backend("torch")
# XXX Temporary fix after https://github.com/pyro-ppl/pyro/pull/2701
import pyro
pyro.enable_validation(False)
encoder = Encoder()
decoder = Decoder()
encode = funsor.function(Reals[28, 28], (Reals[20], Reals[20]))(encoder)
decode = funsor.function(Reals[20], Reals[28, 28])(decoder)
@funsor.interpretation(funsor.montecarlo.MonteCarlo())
def loss_function(data, subsample_scale):
# Lazily sample from the guide.
loc, scale = encode(data)
q = funsor.Independent(dist.Normal(loc['i'], scale['i'], value='z_i'),
'z', 'i', 'z_i')
# Evaluate the model likelihood at the lazy value z.
probs = decode('z')
p = dist.Bernoulli(probs['x', 'y'], value=data['x', 'y'])
p = p.reduce(ops.add, {'x', 'y'})
# Construct an elbo. This is where sampling happens.
elbo = funsor.Integrate(q, p - q, 'z')
elbo = elbo.reduce(ops.add, 'batch') * subsample_scale
loss = -elbo
return loss
train_loader = torch.utils.data.DataLoader(datasets.MNIST(
DATA_PATH, train=True, download=True, transform=transforms.ToTensor()),
batch_size=args.batch_size,
shuffle=True)
encoder.train()
decoder.train()
optimizer = optim.Adam(list(encoder.parameters()) +
list(decoder.parameters()),
lr=1e-3)
for epoch in range(args.num_epochs):
train_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
subsample_scale = float(len(train_loader.dataset) / len(data))
data = data[:, 0, :, :]
data = funsor.Tensor(data, OrderedDict(batch=Bint[len(data)]))
optimizer.zero_grad()
loss = loss_function(data, subsample_scale)
assert isinstance(loss, funsor.Tensor), loss.pretty()
loss.data.backward()
train_loss += loss.item()
optimizer.step()
if batch_idx % 50 == 0:
print(' loss = {}'.format(loss.item()))
if batch_idx and args.smoke_test:
return
print('epoch {} train_loss = {}'.format(epoch, train_loss))
def pytest_runtest_setup(item):
pyro.clear_param_store()
if item.get_marker("disable_validation"):
pyro.enable_validation(False)
else:
pyro.enable_validation(True)
test_initialize_marker = item.get_marker("init")
if test_initialize_marker:
rng_seed = test_initialize_marker.kwargs["rng_seed"]
pyro.set_rng_seed(rng_seed)
def pytest_runtest_setup(item):
pyro.clear_param_store()
if item.get_closest_marker("disable_validation"):
pyro.enable_validation(False)
else:
pyro.enable_validation(True)
test_initialize_marker = item.get_closest_marker("init")
if test_initialize_marker:
rng_seed = test_initialize_marker.kwargs["rng_seed"]
pyro.set_rng_seed(rng_seed)
def forecast(X,y):
# Creating the RBF kernel with the desired lengthscale and variance using pyro.
k1 = gp.kernels.RBF(input_dim=2, lengthscale=torch.tensor(50.0),\
variance = torch.tensor(0.5))
pyro.enable_validation(True)
optim = Adam({"lr": 0.01})
pyro.clear_param_store()
# Creating an array with the last value entered and the 7 weeks ahead (in days).
plus_arr = np.max(X)+np.array([7.,14.,21.,28.,35.,42.,49.,56.,63.,70.])
# Changing numpy arrays into pytorch tensors (Faster for machine learning).
X2 = (torch.from_numpy(X))
y2 = (torch.from_numpy(y-np.mean(y)))
# Adding the new prediction dates into the array and then transforming into pytorch tensor.
Xtest_use = np.append(X,plus_arr)
Xtest_use2 = (torch.from_numpy(Xtest_use))
# Running the GP RBF kernel model on the known data
gpr = gp.models.GPRegression(X2, y2,k1, noise=torch.tensor(0.01))
# Stochastic variational inference to optimise the loss function
# Esentially minimising the errors on the model
svi = SVI(gpr.model, gpr.guide, optim, loss=Trace_ELBO())
losses = []
# Choosing how many times to iterate over the optimisiation
num_steps = 10
for k in range(num_steps):
losses.append(svi.step())
# Putting the results into numpy arrays to be outputted.
with torch.no_grad():
if type(gpr) == gp.models.VariationalSparseGP:
mean, cov = gpr(Xtest_use2, full_cov=True)
else:
mean, cov = gpr(Xtest_use2, full_cov=False, noiseless=False)
mean = mean.detach().numpy()+np.mean(y)
return mean, Xtest_use
def pytest_runtest_setup(item):
np.random.seed(0)
if _BACKEND == "torch":
import pyro
pyro.set_rng_seed(0)
pyro.enable_validation(True)
elif _BACKEND == "jax":
from jax.config import config
config.update('jax_platform_name', 'cpu')
funsor.util.set_backend = _disallow_set_backend
def main(args):
assert pyro.__version__ >= "0.4.1"
pyro.enable_validation(__debug__)
pyro.set_rng_seed(args.seed)
dataset = load_hourly_od(args)
forecaster = train(args, dataset)
num_samples = 10
forecast = forecaster(0, 24 * 7, 24)
assert forecast.shape == (24, ) + dataset["counts"].shape[-2:]
forecast = forecaster(0, 24 * 7, 24, num_samples=num_samples)
assert forecast.shape == (num_samples, 24) + dataset["counts"].shape[-2:]
return forecast
def main(args, reptition=1, path="./IHDP/"):
pyro.enable_validation(__debug__)
# if args.cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
# Generate synthetic data.
pyro.set_rng_seed(args.seed)
train, test, contfeats, binfeats = IHDP(path=path,
reps=reptition,
cuda=True)
(x_train, t_train, y_train), true_ite_train = train
(x_test, t_test, y_test), true_ite_test = test
ym, ys = y_train.mean(), y_train.std()
y_train = (y_train - ym) / ys
# Train.
pyro.set_rng_seed(args.seed)
pyro.clear_param_store()
tedvae = TEDVAE(feature_dim=args.feature_dim,
continuous_dim=contfeats,
binary_dim=binfeats,
latent_dim=args.latent_dim,
latent_dim_t=args.latent_dim_t,
latent_dim_y=args.latent_dim_y,
hidden_dim=args.hidden_dim,
num_layers=args.num_layers,
num_samples=10)
tedvae.fit(x_train,
t_train,
y_train,
num_epochs=args.num_epochs,
batch_size=args.batch_size,
learning_rate=args.learning_rate,
learning_rate_decay=args.learning_rate_decay,
weight_decay=args.weight_decay)
# Evaluate.
est_ite = tedvae.ite(x_test, ym, ys)
est_ite_train = tedvae.ite(x_train, ym, ys)
pehe = np.sqrt(
np.mean((true_ite_test.squeeze() - est_ite.cpu().numpy()) *
(true_ite_test.squeeze() - est_ite.cpu().numpy())))
pehe_train = np.sqrt(
np.mean((true_ite_train.squeeze() - est_ite_train.cpu().numpy()) *
(true_ite_train.squeeze() - est_ite_train.cpu().numpy())))
print("PEHE_train = {:0.3g}".format(pehe_train))
print("PEHE = {:0.3g}".format(pehe))
return pehe, pehe_train
def validate_model(train_loader, encoder=False, transform=False):
pyro.enable_validation(True)
model = ClassifierBnn()
x, y = next(iter(train_loader))
if transform != False:
x = transform(x)
if encoder != False:
z_loc, z_scale = encoder(x)
x = torch.cat((z_loc, z_scale), 1)
print(pyro.poutine.trace(model).get_trace(x, y).format_shapes())
Model = ClassifierBnn()
pyro.enable_validation(True)
pyro.clear_param_store()
model = ClassifierBnn(num_hidden=10, prior_std=1.)
def build(self, input_shape, out, scale_mu = 5.0, scale_lf = 1.0,scale_lv=1.0, df=2,
guide='LaplaceApproximation',
family = "binomial", link = "logit"):
self.create_link(family, link)
#self.link = link
torch.cuda.empty_cache()
pyro.enable_validation(True)
pyro.clear_param_store()
#XX = torch.tensor(X, dtype = torch.float32)
#YY = torch.tensor(Y, dtype = torch.float32)
self.e = input_shape[1]
self.sp = out
self.n = input_shape[0]
self.df = df
e = self.e
sp = self.sp
n = self.n
def model(XX, YY=None, indices=None, posterior = self.posterior):
mu_scale = torch.ones([e, sp])*scale_mu
mu_loc = torch.zeros([e, sp])
lv_scale2 = torch.ones([n, df])*scale_lv
lv_loc = torch.zeros([n, df])
lf_scale2 = torch.ones([df, sp])*scale_lf
lf_loc = torch.zeros([df, sp])
mu = pyro.sample("mu", pyro.distributions.Normal(mu_loc, mu_scale).to_event())
lf = pyro.sample("lf", pyro.distributions.Normal(lf_loc, lf_scale2).to_event())
lv = pyro.sample("lv", pyro.distributions.Normal(lv_loc, lv_scale2).to_event())
with pyro.plate('data', size = XX.shape[0]):
loc_tmp = XX.matmul(mu).add( lv.index_select(0, indices).matmul(lf) )
posterior(loc_tmp, YY)
self.model = model
guide2 = guide
self.elbo = pyro.infer.JitTrace_ELBO(ignore_jit_warnings=True)
if guide2 == 'LaplaceApproximation':
self.guide = pyro.infer.autoguide.AutoLaplaceApproximation(self.model)
if guide2 == 'LowRankMultivariateNormal':
self.guide = pyro.infer.autoguide.AutoLowRankMultivariateNormal(self.model, init_loc_fn=pyro.infer.autoguide.init_to_mean)
self.elbo = pyro.infer.Trace_ELBO(ignore_jit_warnings=True)
if guide2 == "Delta":
self.guide = pyro.infer.autoguide.AutoDelta(self.model)
if guide2 == 'DiagonalNormal':
self.guide = pyro.infer.autoguide.AutoDiagonalNormal(self.model)
def main(args):
pyro.enable_validation(__debug__)
pyro.set_rng_seed(args.rng_seed)
# Generate data.
dataset = generate_data(args)
obs = dataset["obs"]
# Run inference.
model = Model(args, obs)
infer = {"mcmc": infer_mcmc, "svi": infer_svi}[args.infer]
infer(args, model)
# Predict latent time series.
predict(args, model, truth=dataset["S2I"])
def main(args):
pyro.set_rng_seed(0)
pyro.clear_param_store()
pyro.enable_validation(__debug__)
# Loading data
corpora = prepro_file_load("corpora")
documents = list(prepro_file_load("id2pre_text").values())
documents = [re.sub("[\[\]',]", "", doc).split() for doc in documents]
data = [torch.tensor(list(filter(lambda a: a != -1, corpora.doc2idx(doc))), dtype=torch.int64) for doc in documents]
N = list(map(len, data))
args.num_words = len(corpora)
args.num_docs = len(data)
# We'll train using SVI.
logging.info('Training on {} documents'.format(args.num_docs))
predictor = make_predictor(args)
guide = functools.partial(parametrized_guide, predictor)
Elbo = JitTraceEnum_ELBO if args.jit else Trace_ELBO
elbo = Elbo(max_plate_nesting=2)
optim = ClippedAdam({'lr': args.learning_rate})
svi = SVI(model, guide, optim, elbo)
losses = []
logging.info('Step\tLoss')
for step in tqdm(range(args.num_steps)):
loss = svi.step(data, N, args=args)
losses.append(loss)
if step % 10 == 0:
# logging.info('{: >5d}\t{}'.format(step, loss))
logging.info(f"Loss: {loss}")
loss = elbo.loss(model, guide, data, N, args=args)
logging.info('final loss = {}'.format(loss))
# Plot loss over iterations
plt.plot(losses)
plt.title("ELBO")
plt.xlabel("step")
plt.ylabel("loss")
plot_file_name = "../loss-2017_variable-sizes_only-word-data.png"
plt.savefig(plot_file_name)
plt.show()
# save model
torch.save({"model": predictor.state_dict(), "guide": guide}, "../mymodel.pt")
pyro.get_param_store().save("mymodelparams.pt")
def run_svi(itr, x, y):
print(f'Running SVI for {itr} iterations')
pyro.enable_validation(True)
pyro.clear_param_store()
global svi
svi = pyro.infer.SVI(model,
guide,
optim.Adam({"lr": .005}),
loss=pyro.infer.Trace_ELBO(max_plate_nesting=1),
num_samples=1000)
pyro.clear_param_store()
rec_loss = []
for i in range(itr):
loss = svi.step(x, y)
rec_loss.append(loss)
return rec_loss
def infer(self, num_steps=1_000):
model_conditioned = self.model_conditioned()
pyro.clear_param_store()
pyro.enable_validation(True)
svi = pyro.infer.SVI(model=model_conditioned,
guide=self.traits_guide,
optim=pyro.optim.Adam({"lr": 1e-2}),
loss=pyro.infer.TraceGraph_ELBO())
losses = []
for t in tqdm(range(num_steps)):
losses.append(svi.step())
plt.plot(losses)
plt.title("ELBO")
plt.xlabel("step")
plt.ylabel("loss")
def main(args):
logging.info('Generating data')
# WL: edited. =====
# pyro.set_rng_seed(0)
pyro.clear_param_store()
pyro.enable_validation(__debug__)
# ====================
# We can generate synthetic data directly by calling the model.
true_topic_weights, true_topic_words, data = model_original(args=args)
# We'll train using SVI.
logging.info('-' * 40)
logging.info('Training on {} documents'.format(args.num_docs))
# wy: currently don't do enumeration.
# # Elbo = JitTraceEnum_ELBO if args.jit else TraceEnum_ELBO
Elbo = TraceEnum_ELBO
# Elbo = TraceGraph_ELBO
elbo = Elbo(max_plate_nesting=2)
optim = Adam({'lr': args.learning_rate})
svi = SVI(model, guide, optim, elbo)
# WL: edited. =====
# logging.info('Step\tLoss')
logging.info(args)
times = [time.time()]
logging.info('\nstep\t' + 'epoch\t' + 'elbo\t' + 'time(sec)')
# =================
# WL: edited. =====
# for step in range(args.num_steps):
for step in range(1, args.num_steps+1):
# =================
loss = svi.step(data, args=args, batch_size=args.batch_size)
# WL: edited. =====
# if step % 10 == 0:
# logging.info('{: >5d}\t{}'.format(step, loss))
if (step % 10 == 0) or (step == 1):
times.append(time.time())
logging.info(f'{step:06d}\t'
f'{(step * args.batch_size) / args.num_docs:.3f}\t'
f'{-loss:.4f}\t'
f'{times[-1]-times[-2]:.3f}')
def main(args):
pyro.enable_validation(__debug__)
pyro.set_rng_seed(args.rng_seed)
# Generate data.
dataset = generate_data(args)
obs = dataset["obs"]
# Run inference.
model = Model(args, obs)
infer = {"mcmc": infer_mcmc, "svi": infer_svi}[args.infer]
samples = infer(args, model)
# Evaluate fit.
evaluate(args, model, samples)
# Predict latent time series.
if args.forecast:
predict(args, model, truth=dataset["new_I"])
def GP(Y, M, V):
X = np.arange(np.size(M))
k1 = gp.kernels.RBF(input_dim=2, lengthscale=torch.tensor(0.8),\
variance = torch.tensor(2.5))
smoke_test = ('CI' in os.environ
) # ignore; used to check code integrity in the Pyro repo
pyro.enable_validation(True) # can help with debugging
optim = Adam({"lr": 0.01})
pyro.clear_param_store()
plus_arr = np.max(X) + np.array([0.5, 1, 1.5, 2, 2.5])
X2 = (torch.from_numpy(X.astype(float)))
y2 = (torch.from_numpy(V - np.mean(V)))
Xtest_use = np.append(X.astype(float), plus_arr.astype(float))
Xtest_use2 = (torch.from_numpy(Xtest_use))
gpr = gp.models.GPRegression(X2, y2, k1, noise=torch.tensor(0.01))
svi = SVI(gpr.model, gpr.guide, optim, loss=Trace_ELBO())
losses = []
num_steps = 500
for k in range(num_steps):
losses.append(svi.step())
with torch.no_grad():
if type(gpr) == gp.models.VariationalSparseGP:
mean, cov = gpr(Xtest_use2, full_cov=True)
else:
mean, cov = gpr(Xtest_use2, full_cov=False, noiseless=False)
sd = cov.sqrt().detach().numpy()
mean = mean.detach().numpy() + np.mean(V)
#for param_name in pyro.get_param_store().get_all_param_names():
# print('{}={}'.format(param_name,pyro.param(param_name).item()))
return mean, Xtest_use, X
def build_model(self, X, Y,Re,SP, df, scale_mu = 5.0, scale_scale = 1.0, batch_size = 20):
pyro.enable_validation(True)
pyro.clear_param_store()
#XX = torch.tensor(X, dtype = torch.float32)
#YY = torch.tensor(Y, dtype = torch.float32)
self.e = X.shape[1]
self.sp = Y.shape[1]
self.df = df
e = self.e
sp = self.sp
if Re is None:
def model(XX, YY=None, link = self.link):
mu_scale = torch.ones([e, sp])
mu_loc = torch.zeros([e, sp])
scale_scale2 = torch.ones([sp, df])
scale_loc = torch.zeros([sp, df])
mu = pyro.sample("mu", pyro.distributions.Normal(mu_loc, mu_scale).to_event())
scale = pyro.sample("scale", pyro.distributions.Normal(scale_loc, scale_scale2).to_event())
with pyro.plate('data', size = XX.shape[0]):
loc_tmp = XX.matmul(mu)
#loc_tmp = mu(XX[ind,:])
pyro.sample("obs", MultivariateProbit(loc_tmp, scale), obs = YY)
else:
len_unique = np.unique(Re[:,0]).shape[0]
#indices = torch.tensor(Re, dtype=torch.long)
def model(XX, YY=None, Re=None, link = self.link):
mu_scale = torch.ones([e, sp])*scale_mu
mu_loc = torch.zeros([e, sp])
scale_scale2 = torch.ones([sp, df])*scale_scale
scale_loc = torch.zeros([sp, df])
mu = pyro.sample("mu", pyro.distributions.Normal(mu_loc, mu_scale).to_event())
scale = pyro.sample("scale", pyro.distributions.Normal(scale_loc, scale_scale2).to_event())
re = pyro.sample("re", pyro.distributions.Normal(0.0, torch.ones([len_unique,1])).to_event())
with pyro.plate('data', size = XX.shape[0]):
if Re is None:
loc_tmp = XX.matmul(mu)
else:
loc_tmp = XX.matmul(mu) + re.gather(0, Re)
pyro.sample("obs", MultivariateProbit(loc_tmp, scale), obs = YY)
return model
def main(args):
pyro.set_rng_seed(0)
pyro.enable_validation()
optim = Adam({"lr": 0.1})
inference = SVI(model, guide, optim, loss=Trace_ELBO())
data = torch.tensor([0.0, 1.0, 2.0, 20.0, 30.0, 40.0])
k = 2
print('Step\tLoss')
loss = 0.0
for step in range(args.num_epochs):
if step and step % 10 == 0:
print('{}\t{:0.5g}'.format(step, loss))
loss = 0.0
loss += inference.step(data, k)
print('Parameters:')
for name in sorted(pyro.get_param_store().get_all_param_names()):
print('{} = {}'.format(name, pyro.param(name).detach().cpu().numpy()))
def train(prior, data, num_steps, num_obs):
pyro.enable_validation(True)
pyro.clear_param_store()
# max_plate_nesting = 1 because there is a single plate in the model
loss_func = pyro.infer.TraceEnum_ELBO(max_plate_nesting=1)
svi = pyro.infer.SVI(model,
guide,
pyro.optim.Adam({'lr': .01}),
loss=loss_func
)
losses = []
for _ in tqdm(range(num_steps)):
loss = svi.step(prior, data, num_obs)
losses.append(loss)
plt.figure()
plt.plot(losses)
plt.show()
posterior_params = {k: np.array(v.data) for k, v in pyro.get_param_store().items()}
posterior_params['a'] = posterior_params['a'][None, :] # reshape to same as other variables
return posterior_params
def main(args):
pyro.set_rng_seed(0)
pyro.enable_validation()
optim = Adam({"lr": 0.1})
inference = SVI(model, guide, optim, loss=Trace_ELBO())
data = torch.tensor([0.0, 1.0, 2.0, 20.0, 30.0, 40.0])
k = 2
print('Step\tLoss')
loss = 0.0
for step in range(args.num_epochs):
if step and step % 10 == 0:
print('{}\t{:0.5g}'.format(step, loss))
loss = 0.0
loss += inference.step(data, k)
print('Parameters:')
for name in sorted(pyro.get_param_store().get_all_param_names()):
print('{} = {}'.format(name, pyro.param(name).detach().cpu().numpy()))
def model_gamma(X, y, column_names):
pyro.enable_validation(True)
min_value = torch.finfo(X.dtype).eps
max_value = torch.finfo(X.dtype).max
# We still need to calculate our linear combination
intercept_prior = dist.Normal(0.0, 1.0)
linear_combination = pyro.sample(f"beta_intercept", intercept_prior)
#print("intercept", linear_combination)
# Also define coefficient priors
for i in range(X.shape[1]):
coefficient_prior = dist.Normal(0.0, 1.0)
beta_coef = pyro.sample(f"beta_{column_names[i]}", coefficient_prior)
#print(column_names[i], beta_coef)
linear_combination = linear_combination + (X[:, i] * beta_coef)
# But now our mean will be e^{linear combination}
mean = torch.exp(linear_combination).clamp(min=min_value, max=max_value)
# We will also define a rate parameter
rate = pyro.sample("rate", dist.HalfNormal(scale=10.0)).clamp(min=min_value)
# Since mean = shape/rate, then the shape = mean * rate
shape = (mean * rate)
# Now that we have the shape and rate parameters for the
# Gamma distribution, we can draw samples from it and condition
# them on our observations
with pyro.plate("data", y.shape[0]):
outcome_dist = dist.Gamma(shape, rate)
observation = pyro.sample("obs", outcome_dist, obs=y)
the No U-Turn Sampler, which provides an efficient and automated way (i.e. limited
hyper-parameters) of running HMC on different problems.
[1] Carpenter B. (2016), ["Hierarchical Partial Pooling for Repeated Binary Trials"]
(http://mc-stan.org/users/documentation/case-studies/pool-binary-trials.html).
[2] Efron B., Morris C. (1975), "Data analysis using Stein's estimator and its
generalizations", J. Amer. Statist. Assoc., 70, 311-319.
[3] Neal, R. (2012), "MCMC using Hamiltonian Dynamics",
(https://arxiv.org/pdf/1206.1901.pdf)
[4] Hoffman, M. D. and Gelman, A. (2014), "The No-U-turn sampler: Adaptively setting
path lengths in Hamiltonian Monte Carlo", (https://arxiv.org/abs/1111.4246)
"""
logging.basicConfig(format='%(message)s', level=logging.INFO)
# Enable validation checks
pyro.enable_validation(True)
pyro.set_rng_seed(1)
DATA_URL = "https://d2fefpcigoriu7.cloudfront.net/datasets/EfronMorrisBB.txt"
# ===================================
# MODELS
# ===================================
def fully_pooled(at_bats):
"""
Number of hits in $K$ at bats for each player has a Binomial
distribution with a common probability of success, $\phi$.
:param (torch.Tensor) at_bats: Number of at bats for each player.
