def pytest_configure(config): try: import funsor except ImportError: pass else: funsor.set_backend("torch")
def __init__(self, *args, **kwargs): if collapse._coerce is None: import funsor from funsor.distribution import CoerceDistributionToFunsor funsor.set_backend("jax") collapse._coerce = CoerceDistributionToFunsor("jax") super().__init__(*args, **kwargs)
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 __init__(self, *args, **kwargs): if CollapseMessenger._coerce is None: import funsor from funsor.distribution import CoerceDistributionToFunsor funsor.set_backend("torch") CollapseMessenger._coerce = CoerceDistributionToFunsor("torch") self._block = False super().__init__(*args, **kwargs)
def _import_funsor(): try: import funsor except ImportError as e: raise ImportError( 'AutoGaussian(..., backend="funsor") requires funsor. ' "Try installing via: pip install pyro-ppl[funsor]" ) from e funsor.set_backend("torch") return funsor
def main(args): funsor.set_backend("torch") # Declare parameters. trans_probs = torch.tensor([[0.2, 0.8], [0.7, 0.3]], requires_grad=True) emit_probs = torch.tensor([[0.4, 0.6], [0.1, 0.9]], requires_grad=True) params = [trans_probs, emit_probs] # A discrete HMM model. def model(data): log_prob = funsor.to_funsor(0.) trans = dist.Categorical(probs=funsor.Tensor( trans_probs, inputs=OrderedDict([('prev', funsor.Bint[args.hidden_dim])]), )) emit = dist.Categorical(probs=funsor.Tensor( emit_probs, inputs=OrderedDict([('latent', funsor.Bint[args.hidden_dim])]), )) x_curr = funsor.Number(0, args.hidden_dim) for t, y in enumerate(data): x_prev = x_curr # A delayed sample statement. x_curr = funsor.Variable('x_{}'.format(t), funsor.Bint[args.hidden_dim]) log_prob += trans(prev=x_prev, value=x_curr) if not args.lazy and isinstance(x_prev, funsor.Variable): log_prob = log_prob.reduce(ops.logaddexp, x_prev.name) log_prob += emit(latent=x_curr, value=funsor.Tensor(y, dtype=2)) log_prob = log_prob.reduce(ops.logaddexp) return log_prob # Train model parameters. data = torch.ones(args.time_steps, dtype=torch.long) optim = torch.optim.Adam(params, lr=args.learning_rate) for step in range(args.train_steps): optim.zero_grad() if args.lazy: with interpretation(lazy): log_prob = apply_optimizer(model(data)) log_prob = reinterpret(log_prob) else: log_prob = model(data) assert not log_prob.inputs, 'free variables remain' loss = -log_prob.data loss.backward() optim.step()
def main(args): funsor.set_backend("torch") # Define a basic model with a single Normal latent random variable `loc` # and a batch of Normally distributed observations. def model(data): loc = pyro.sample("loc", dist.Normal(0., 1.)) with pyro.plate("data", len(data), dim=-1): pyro.sample("obs", dist.Normal(loc, 1.), obs=data) # Define a guide (i.e. variational distribution) with a Normal # distribution over the latent random variable `loc`. def guide(data): guide_loc = pyro.param("guide_loc", torch.tensor(0.)) guide_scale = pyro.param("guide_scale", torch.tensor(1.), constraint=constraints.positive) pyro.sample("loc", dist.Normal(guide_loc, guide_scale)) # Generate some data. torch.manual_seed(0) data = torch.randn(100) + 3.0 # Because the API in minipyro matches that of Pyro proper, # training code works with generic Pyro implementations. with pyro_backend(args.backend), interpretation(MonteCarlo()): # Construct an SVI object so we can do variational inference on our # model/guide pair. Elbo = infer.JitTrace_ELBO if args.jit else infer.Trace_ELBO elbo = Elbo() adam = optim.Adam({"lr": args.learning_rate}) svi = infer.SVI(model, guide, adam, elbo) # Basic training loop pyro.get_param_store().clear() for step in range(args.num_steps): loss = svi.step(data) if args.verbose and step % 100 == 0: print("step {} loss = {}".format(step, loss)) # Report the final values of the variational parameters # in the guide after training. if args.verbose: for name in pyro.get_param_store(): value = pyro.param(name).data print("{} = {}".format(name, value.detach().cpu().numpy())) # For this simple (conjugate) model we know the exact posterior. In # particular we know that the variational distribution should be # centered near 3.0. So let's check this explicitly. assert (pyro.param("guide_loc") - 3.0).abs() < 0.1
def main(args): funsor.set_backend("torch") # Declare parameters. trans_noise = torch.tensor(0.1, requires_grad=True) emit_noise = torch.tensor(0.5, requires_grad=True) params = [trans_noise, emit_noise] # A Gaussian HMM model. def model(data): log_prob = funsor.to_funsor(0.) x_curr = funsor.Tensor(torch.tensor(0.)) for t, y in enumerate(data): x_prev = x_curr # A delayed sample statement. x_curr = funsor.Variable('x_{}'.format(t), funsor.Real) log_prob += dist.Normal(1 + x_prev / 2., trans_noise, value=x_curr) # Optionally marginalize out the previous state. if t > 0 and not args.lazy: log_prob = log_prob.reduce(ops.logaddexp, x_prev.name) # An observe statement. log_prob += dist.Normal(0.5 + 3 * x_curr, emit_noise, value=y) # Marginalize out all remaining delayed variables. log_prob = log_prob.reduce(ops.logaddexp) return log_prob # Train model parameters. torch.manual_seed(0) data = torch.randn(args.time_steps) optim = torch.optim.Adam(params, lr=args.learning_rate) for step in range(args.train_steps): optim.zero_grad() if args.lazy: with interpretation(lazy): log_prob = apply_optimizer(model(data)) log_prob = reinterpret(log_prob) else: log_prob = model(data) assert not log_prob.inputs, 'free variables remain' loss = -log_prob.data loss.backward() optim.step() if args.verbose and step % 10 == 0: print('step {} loss = {}'.format(step, loss.item()))
def main(args): funsor.set_backend("torch") if args.force or not args.metrics_filename or not os.path.exists( args.metrics_filename): results = track(args) else: results = torch.load(args.metrics_filename) if args.plot_filename: import matplotlib matplotlib.use('Agg') from matplotlib import pyplot import numpy as np seeds = set(seed for seed, _, _ in results) X = args.num_frames pyplot.figure(figsize=(5, 1.4), dpi=300) pos_error = np.array( [[results[s, 0, f]['final_pos_error'] for s in seeds] for f in args.num_frames]) mse = (pos_error**2).mean(axis=1) std = (pos_error**2).std(axis=1) / len(seeds)**0.5 pyplot.plot(X, mse**0.5, 'k--') pyplot.fill_between(X, (mse - std)**0.5, (mse + std)**0.5, color='black', alpha=0.15, lw=0) pos_error = np.array( [[results[s, 1, f]['final_pos_error'] for s in seeds] for f in args.num_frames]) mse = (pos_error**2).mean(axis=1) std = (pos_error**2).std(axis=1) / len(seeds)**0.5 pyplot.plot(X, mse**0.5, 'r-') pyplot.fill_between(X, (mse - std)**0.5, (mse + std)**0.5, color='red', alpha=0.15, lw=0) pyplot.ylabel('Position RMSE') pyplot.xlabel('Track Length') pyplot.xticks((5, 10, 15, 20, 25, 30)) pyplot.xlim(5, 30) pyplot.tight_layout(0) pyplot.savefig(args.plot_filename)
def main(args): funsor.set_backend("torch") torch.manual_seed(args.seed) print_ = print if args.verbose else lambda msg: None print_('Data:') data = torch.distributions.Categorical(torch.ones(2)).sample((args.size, )) assert data.shape == (args.size, ) data = Tensor(data, OrderedDict(i=Bint[args.size]), dtype=2) print_(data) print_('Model:') m = model(args.size) print_(m.pretty()) print_('Eager log_prob:') obs = {str(i): data(i) for i in range(args.size)} log_prob = m(**obs) print_(log_prob)
def test_gaussian_funsor(batch_shape): # This tests sample distribution, rsample gradients, log_prob, and log_prob # gradients for both Pyro's and Funsor's Gaussian. import funsor funsor.set_backend("torch") num_samples = 100000 # Declare unconstrained parameters. loc = torch.randn(batch_shape + (3, )).requires_grad_() t = transform_to(constraints.positive_definite) m = torch.randn(batch_shape + (3, 3)) precision_unconstrained = t.inv(m @ m.transpose(-1, -2)).requires_grad_() # Transform to constrained space. log_normalizer = torch.zeros(batch_shape) precision = t(precision_unconstrained) info_vec = (precision @ loc[..., None])[..., 0] def check_equal(actual, expected, atol=0.01, rtol=0): assert_close(actual.data, expected.data, atol=atol, rtol=rtol) grads = torch.autograd.grad( (actual - expected).abs().sum(), [loc, precision_unconstrained], retain_graph=True, ) for grad in grads: assert grad.abs().max() < atol entropy = dist.MultivariateNormal(loc, precision_matrix=precision).entropy() # Monte carlo estimate entropy via pyro. p_gaussian = Gaussian(log_normalizer, info_vec, precision) p_log_Z = p_gaussian.event_logsumexp() p_rsamples = p_gaussian.rsample((num_samples, )) pp_entropy = (p_log_Z - p_gaussian.log_density(p_rsamples)).mean(0) check_equal(pp_entropy, entropy) # Monte carlo estimate entropy via funsor. inputs = OrderedDict([(k, funsor.Bint[v]) for k, v in zip("ij", batch_shape)]) inputs["x"] = funsor.Reals[3] f_gaussian = funsor.gaussian.Gaussian(mean=loc, precision=precision, inputs=inputs) f_log_Z = f_gaussian.reduce(funsor.ops.logaddexp, "x") sample_inputs = OrderedDict(particle=funsor.Bint[num_samples]) deltas = f_gaussian.sample("x", sample_inputs) f_rsamples = funsor.montecarlo.extract_samples(deltas)["x"] ff_entropy = (f_log_Z - f_gaussian(x=f_rsamples)).reduce( funsor.ops.mean, "particle") check_equal(ff_entropy.data, entropy) # Check Funsor's .rsample against Pyro's .log_prob. pf_entropy = (p_log_Z - p_gaussian.log_density(f_rsamples.data)).mean(0) check_equal(pf_entropy, entropy) # Check Pyro's .rsample against Funsor's .log_prob. fp_rsamples = funsor.Tensor(p_rsamples)["particle"] for i in "ij"[:len(batch_shape)]: fp_rsamples = fp_rsamples[i] fp_entropy = (f_log_Z - f_gaussian(x=fp_rsamples)).reduce( funsor.ops.mean, "particle") check_equal(fp_entropy.data, entropy)
def stats( model: avail_models = typer.Option("cosmos", help="Tapqir model", prompt="Tapqir model"), cuda: bool = typer.Option( partial(get_default, "cuda"), "--cuda/--cpu", help="Run computations on GPU or CPU", prompt="Run computations on GPU?", show_default=False, ), nbatch_size: int = typer.Option( partial(get_default, "nbatch-size"), "--nbatch-size", "-n", help="AOI batch size", prompt="AOI batch size", ), fbatch_size: int = typer.Option( partial(get_default, "fbatch-size"), "--fbatch-size", "-f", help="Frame batch size", prompt="Frame batch size", ), matlab: bool = typer.Option( partial(get_default, "matlab"), "--matlab", help="Save parameters in matlab format", prompt="Save parameters in matlab format?", ), funsor: bool = typer.Option(False, "--funsor/--pyro", help="Use funsor or pyro backend"), no_input: bool = typer.Option( False, "--no-input", help="Use defaults values.", is_eager=True, callback=deactivate_prompts, ), ): from pyroapi import pyro_backend from tapqir.models import models logger = logging.getLogger("tapqir") global DEFAULTS cd = DEFAULTS["cd"] dtype = "double" device = "cuda" if cuda else "cpu" backend = "funsor" if funsor else "pyro" settings = {} settings["device"] = device settings["dtype"] = dtype # pyro backend if backend == "pyro": PYRO_BACKEND = "pyro" elif backend == "funsor": import funsor import pyro.contrib.funsor # noqa: F401 funsor.set_backend("torch") PYRO_BACKEND = "contrib.funsor" else: raise ValueError("Only pyro and funsor backends are supported.") with pyro_backend(PYRO_BACKEND): logger.info("Computing stats ...") model = models[model](**settings) try: model.load(cd) except TapqirFileNotFoundError: logger.exception("Failed to load data file") return 1 model.load_checkpoint(param_only=True) model.nbatch_size = nbatch_size model.fbatch_size = fbatch_size try: model.compute_stats(save_matlab=matlab) except CudaOutOfMemoryError: logger.exception("Failed to compute stats") return 1 logger.info("Computing stats: Done") return 0
def fit( model: avail_models = typer.Option( "cosmos", help="Tapqir model", prompt="Tapqir model"), cuda: bool = typer.Option( partial(get_default, "cuda"), "--cuda/--cpu", help="Run computations on GPU or CPU", prompt="Run computations on GPU?", show_default=False, ), nbatch_size: int = typer.Option( partial(get_default, "nbatch-size"), "--nbatch-size", "-n", help="AOI batch size", prompt="AOI batch size", ), fbatch_size: int = typer.Option( partial(get_default, "fbatch-size"), "--fbatch-size", "-f", help="Frame batch size", prompt="Frame batch size", ), learning_rate: float = typer.Option( partial(get_default, "learning-rate"), "--learning-rate", "-lr", help="Learning rate", prompt="Learning rate", ), num_iter: int = typer.Option( 0, "--num-iter", "-it", help="Number of iterations", prompt="Number of iterations", ), k_max: int = typer.Option( 2, "--k-max", "-k", help="Maximum number of spots per image"), matlab: bool = typer.Option( partial(get_default, "matlab"), "--matlab", help="Save parameters in matlab format", prompt="Save parameters in matlab format?", ), funsor: bool = typer.Option( False, "--funsor/--pyro", help="Use funsor or pyro backend"), pykeops: bool = typer.Option( True, "--pykeops/--no-pykeops", help="Use pykeops backend for offset marginalization", ), overwrite: bool = typer.Option( True, "--overwrite", "-w", help="Overwrite defaults values.", prompt="Overwrite defaults values?", ), no_input: bool = typer.Option( False, "--no-input", help="Disable interactive prompt.", is_eager=True, callback=deactivate_prompts, ), progress_bar=None, ): """ Fit the data to the selected model. Available models: * cosmos: single-color time-independent co-localization model.\n """ global DEFAULTS cd = DEFAULTS["cd"] if progress_bar is None: progress_bar = tqdm from pyroapi import pyro_backend from tapqir.models import models logger = logging.getLogger("tapqir") settings = {} settings["K"] = k_max settings["device"] = "cuda" if cuda else "cpu" settings["dtype"] = "double" settings["use_pykeops"] = pykeops # priors settings settings["priors"] = DEFAULTS["priors"] if overwrite: DEFAULTS["cuda"] = cuda DEFAULTS["nbatch-size"] = nbatch_size DEFAULTS["fbatch-size"] = fbatch_size DEFAULTS["learning-rate"] = learning_rate DEFAULTS["matlab"] = matlab with open(cd / ".tapqir" / "config.yaml", "w") as cfg_file: yaml.dump( {key: value for key, value in DEFAULTS.items() if key != "cd"}, cfg_file, sort_keys=False, ) backend = "funsor" if funsor else "pyro" if model == "cosmos+hmm": backend = "funsor" # hmm requires funsor backend if backend == "pyro": PYRO_BACKEND = "pyro" elif backend == "funsor": import funsor import pyro.contrib.funsor # noqa: F401 funsor.set_backend("torch") PYRO_BACKEND = "contrib.funsor" else: raise ValueError("Only pyro and funsor backends are supported.") with pyro_backend(PYRO_BACKEND): logger.info("Fitting the data ...") model = models[model](**settings) try: model.load(cd) except TapqirFileNotFoundError as err: logger.exception(f"Failed to load {err.name} file") return 1 model.init(learning_rate, nbatch_size, fbatch_size) try: model.run(num_iter, progress_bar=progress_bar) except CudaOutOfMemoryError: logger.exception("Failed to fit the data") return 1 logger.info("Fitting the data: Done") logger.info("Computing stats ...") try: model.compute_stats(save_matlab=matlab) except CudaOutOfMemoryError: logger.exception("Failed to compute stats") return 1 logger.info("Computing stats: Done") return 0
def stats( model: Model = typer.Option("cosmos", help="Tapqir model", prompt="Tapqir model"), channels: List[int] = typer.Option( [0], help="Color-channel numbers to analyze", prompt="Channel numbers (space separated if multiple)", ), cuda: bool = typer.Option( partial(get_default, "cuda"), "--cuda/--cpu", help="Run computations on GPU or CPU", prompt="Run computations on GPU?", show_default=False, ), nbatch_size: int = typer.Option( partial(get_default, "nbatch-size"), "--nbatch-size", "-n", help="AOI batch size", prompt="AOI batch size", ), fbatch_size: int = typer.Option( partial(get_default, "fbatch-size"), "--fbatch-size", "-f", help="Frame batch size", prompt="Frame batch size", ), matlab: bool = typer.Option( partial(get_default, "matlab"), "--matlab", help="Save parameters in matlab format", prompt="Save parameters in matlab format?", ), funsor: bool = typer.Option( False, "--funsor/--pyro", help="Use funsor or pyro backend" ), no_input: bool = typer.Option( False, "--no-input", help="Use defaults values.", is_eager=True, callback=deactivate_prompts, ), ): from pyroapi import pyro_backend from tapqir.models import models from tapqir.utils.stats import save_stats global DEFAULTS cd = DEFAULTS["cd"] dtype = "double" device = "cuda" if cuda else "cpu" backend = "funsor" if funsor else "pyro" # pyro backend if backend == "pyro": PYRO_BACKEND = "pyro" elif backend == "funsor": import funsor import pyro.contrib.funsor # noqa: F401 funsor.set_backend("torch") PYRO_BACKEND = "contrib.funsor" else: raise ValueError("Only pyro and funsor backends are supported.") with pyro_backend(PYRO_BACKEND): model = models[model](1, 2, channels, device, dtype) model.load(cd) model.load_checkpoint(param_only=True) model.nbatch_size = nbatch_size model.fbatch_size = fbatch_size typer.echo("Computing stats ...") save_stats(model, cd, save_matlab=matlab) typer.echo("Computing stats: Done")
# Copyright Contributors to the Pyro project. # SPDX-License-Identifier: Apache-2.0 try: import funsor except ImportError: raise ImportError( "Looking like you want to do inference for models with " "discrete latent variables. This is an experimental feature. " "You need to install `funsor` to be able to use this feature. " "It can be installed with `pip install funsor`.") from numpyro.contrib.funsor.enum_messenger import enum, infer_config, markov, plate, to_data, to_funsor, trace from numpyro.contrib.funsor.infer_util import config_enumerate, log_density, plate_to_enum_plate funsor.set_backend("jax") __all__ = [ "config_enumerate", "enum", "infer_config", "log_density", "markov", "plate", "plate_to_enum_plate", "to_data", "to_funsor", "trace", ]
def main(args): funsor.set_backend("torch") # download and pre-process EEG data if not in test mode if not args.test: download_data() N_val, N_test = 149, 200 data = np.loadtxt('eeg.dat', delimiter=',', skiprows=19) print("[raw data shape] {}".format(data.shape)) data = data[::20, :] print("[data shape after thinning] {}".format(data.shape)) eye_state = [int(d) for d in data[:, -1].tolist()] data = torch.tensor(data[:, :-1]).float() # in test mode (for continuous integration on github) so create fake data else: data = torch.randn(10, 3) N_val, N_test = 2, 2 T, obs_dim = data.shape N_train = T - N_test - N_val np.random.seed(0) rand_perm = np.random.permutation(N_val + N_test) val_indices = rand_perm[0:N_val] test_indices = rand_perm[N_val:] data_mean = data[0:N_train, :].mean(0) data -= data_mean data_std = data[0:N_train, :].std(0) data /= data_std print("Length of time series T: {} Observation dimension: {}".format(T, obs_dim)) print("N_train: {} N_val: {} N_test: {}".format(N_train, N_val, N_test)) torch.manual_seed(args.seed) # set up model slds = SLDS(num_components=args.num_components, hidden_dim=args.hidden_dim, obs_dim=obs_dim, fine_observation_noise=args.fon, fine_transition_noise=args.ftn, fine_observation_matrix=args.fom, fine_transition_matrix=args.ftm, moment_matching_lag=args.moment_matching_lag) # set up optimizer adam = torch.optim.Adam(slds.parameters(), lr=args.learning_rate, betas=(args.beta1, 0.999), amsgrad=True) scheduler = torch.optim.lr_scheduler.ExponentialLR(adam, gamma=args.gamma) ts = [time.time()] report_frequency = 1 # training loop for step in range(args.num_steps): nll = -slds.log_prob(data[0:N_train, :]) / N_train nll.backward() if step == 5: scheduler.base_lrs[0] *= 0.20 adam.step() scheduler.step() adam.zero_grad() if step % report_frequency == 0 or step == args.num_steps - 1: step_dt = ts[-1] - ts[-2] if step > 0 else 0.0 pred_mse, pred_LLs = slds.filter_and_predict(data[0:N_train + N_val + N_test, :]) val_mse = pred_mse[val_indices].mean().item() test_mse = pred_mse[test_indices].mean().item() val_ll = pred_LLs[val_indices].mean().item() test_ll = pred_LLs[test_indices].mean().item() stats = "[step %03d] train_nll: %.5f val_mse: %.5f val_ll: %.5f test_mse: %.5f test_ll: %.5f\t(dt: %.2f)" print(stats % (step, nll.item(), val_mse, val_ll, test_mse, test_ll, step_dt)) ts.append(time.time()) # plot predictions and smoothed means if args.plot: assert not args.test predicted_mse, LLs, pred_means, pred_vars, smooth_means, smooth_probs = \ slds.filter_and_predict(data, smoothing=True) pred_means = pred_means.data.numpy() pred_stds = pred_vars.sqrt().data.numpy() smooth_means = smooth_means.data.numpy() smooth_probs = smooth_probs.data.numpy() import matplotlib matplotlib.use('Agg') # noqa: E402 import matplotlib.pyplot as plt f, axes = plt.subplots(4, 1, figsize=(12, 8), sharex=True) T = data.size(0) N_valtest = N_val + N_test to_seconds = 117.0 / T for k, ax in enumerate(axes[:-1]): which = [0, 4, 10][k] ax.plot(to_seconds * np.arange(T), data[:, which], 'ko', markersize=2) ax.plot(to_seconds * np.arange(N_train), smooth_means[:N_train, which], ls='solid', color='r') ax.plot(to_seconds * (N_train + np.arange(N_valtest)), pred_means[-N_valtest:, which], ls='solid', color='b') ax.fill_between(to_seconds * (N_train + np.arange(N_valtest)), pred_means[-N_valtest:, which] - 1.645 * pred_stds[-N_valtest:, which], pred_means[-N_valtest:, which] + 1.645 * pred_stds[-N_valtest:, which], color='lightblue') ax.set_ylabel("$y_{%d}$" % (which + 1), fontsize=20) ax.tick_params(axis='both', which='major', labelsize=14) axes[-1].plot(to_seconds * np.arange(T), eye_state, 'k', ls='solid') axes[-1].plot(to_seconds * np.arange(T), smooth_probs, 'r', ls='solid') axes[-1].set_xlabel("Time (s)", fontsize=20) axes[-1].set_ylabel("Eye state", fontsize=20) axes[-1].tick_params(axis='both', which='major', labelsize=14) plt.tight_layout(pad=0.7) plt.savefig('eeg.pdf')
def run_expt(args): funsor.set_backend("torch") optim = args["optim"] lr = args["learnrate"] schedule = [] if not args["schedule"] else [ int(i) for i in args["schedule"].split(",") ] # default these to "none" instead of None, which argparse does for some reason args["group"] = "none" if args["group"] is None else args["group"] args["individual"] = "none" if args["individual"] is None else args[ "individual"] random_effects = {"group": args["group"], "individual": args["individual"]} pyro.enable_validation(args["validation"]) pyro.set_rng_seed( args["seed"]) # reproducible random effect parameter init if args["cuda"]: torch.set_default_tensor_type(torch.cuda.FloatTensor) if args["dataset"] == "seal": filename = os.path.join(args["folder"], "prep_seal_data.csv") config = prepare_seal(filename, random_effects) elif args["dataset"] == "fake": fake_sizes = { "state": args["size_state"], "random": args["size_random"], "group": args["size_group"], "individual": args["size_individual"], "timesteps": args["size_timesteps"], } config = prepare_fake(fake_sizes, random_effects) else: raise ValueError("Dataset {} not yet included".format(args["dataset"])) if args["smoke"]: args["timesteps"] = 2 config["sizes"]["timesteps"] = 3 if args["truncate"] > 0: config["sizes"]["timesteps"] = args["truncate"] config["zeroinflation"] = args["zeroinflation"] model = Model(config) guide = Guide(config) loss_fn = parallel_loss_fn if args["jit"]: loss_fn = torch.jit.trace( lambda: loss_fn(model, guide, args["parallel"]), ()) else: loss_fn = functools.partial(loss_fn, model, guide, args["parallel"]) # count the number of parameters once num_parameters = aic_num_parameters(model, guide) losses = [] # TODO support continuous random effects with monte carlo assert random_effects["group"] != "continuous" assert random_effects["individual"] != "continuous" with pyro.poutine.trace(param_only=True) as param_capture: loss_fn() params = [ site["value"].unconstrained() for site in param_capture.trace.nodes.values() ] if optim == "sgd": optimizer = torch.optim.Adam(params, lr=lr) elif optim == "lbfgs": optimizer = torch.optim.LBFGS(params, lr=lr) else: raise ValueError("{} not supported optimizer".format(optim)) if schedule: scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=schedule, gamma=0.5) schedule_step_loss = False else: scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, 'min') schedule_step_loss = True for t in range(args["timesteps"]): def closure(): optimizer.zero_grad() loss = loss_fn() loss.backward() return loss loss = optimizer.step(closure) scheduler.step(loss.item() if schedule_step_loss else t) losses.append(loss.item()) print("Loss: {}, AIC[{}]: ".format(loss.item(), t), 2. * loss + 2. * num_parameters) aic_final = 2. * losses[-1] + 2. * num_parameters print("AIC final: {}".format(aic_final)) results = {} results["args"] = args results["sizes"] = config["sizes"] results["likelihoods"] = losses results["likelihood_final"] = losses[-1] results["aic_final"] = aic_final results["aic_num_parameters"] = num_parameters if args["resultsdir"] is not None and os.path.exists(args["resultsdir"]): re_str = "g" + ("n" if args["group"] is None else "d" if args["group"] == "discrete" else "c") re_str += "i" + ("n" if args["individual"] is None else "d" if args["individual"] == "discrete" else "c") results_filename = "expt_{}_{}_{}.json".format( args["dataset"], re_str, str(uuid.uuid4().hex)[0:5]) with open(os.path.join(args["resultsdir"], results_filename), "w") as f: json.dump(results, f) return results
# SPDX-License-Identifier: Apache-2.0 import logging import pytest import torch from pyro.ops.indexing import Vindex from tests.common import xfail_param # put all funsor-related imports here, so test collection works without funsor try: import funsor import pyro.contrib.funsor funsor.set_backend("torch") from pyroapi import distributions as dist from pyroapi import infer, pyro from tests.contrib.funsor.test_valid_models_enum import assert_ok except ImportError: pytestmark = pytest.mark.skip(reason="funsor is not installed") logger = logging.getLogger(__name__) @pytest.mark.parametrize('enumerate_', [None, "parallel", "sequential"]) def test_enum_discrete_non_enumerated_plate_ok(enumerate_): def model(): pyro.sample("w", dist.Bernoulli(0.5), infer={'enumerate': 'parallel'})
def fit( model: Model = typer.Option("cosmos", help="Tapqir model", prompt="Tapqir model"), channels: List[int] = typer.Option( [0], help="Color-channel numbers to analyze", prompt="Channel numbers (space separated if multiple)", ), cuda: bool = typer.Option( partial(get_default, "cuda"), "--cuda/--cpu", help="Run computations on GPU or CPU", prompt="Run computations on GPU?", show_default=False, ), nbatch_size: int = typer.Option( partial(get_default, "nbatch-size"), "--nbatch-size", "-n", help="AOI batch size", prompt="AOI batch size", ), fbatch_size: int = typer.Option( partial(get_default, "fbatch-size"), "--fbatch-size", "-f", help="Frame batch size", prompt="Frame batch size", ), learning_rate: float = typer.Option( partial(get_default, "learning-rate"), "--learning-rate", "-lr", help="Learning rate", prompt="Learning rate", ), num_iter: int = typer.Option( 0, "--num-iter", "-it", help="Number of iterations", prompt="Number of iterations", ), k_max: int = typer.Option( 2, "--k-max", "-k", help="Maximum number of spots per image" ), matlab: bool = typer.Option( partial(get_default, "matlab"), "--matlab", help="Save parameters in matlab format", prompt="Save parameters in matlab format?", ), funsor: bool = typer.Option( False, "--funsor/--pyro", help="Use funsor or pyro backend" ), pykeops: bool = typer.Option( True, "--pykeops/--no-pykeops", help="Use pykeops backend for offset marginalization", ), overwrite: bool = typer.Option( True, "--overwrite", "-w", help="Overwrite defaults values.", prompt="Overwrite defaults values?", ), no_input: bool = typer.Option( False, "--no-input", help="Disable interactive prompt.", is_eager=True, callback=deactivate_prompts, ), progress_bar=None, ): """ Fit the data to the selected model. Available models: * cosmos: single-color time-independent co-localization model.\n """ global DEFAULTS cd = DEFAULTS["cd"] if progress_bar is None: progress_bar = tqdm from pyroapi import pyro_backend from tapqir.models import models from tapqir.utils.stats import save_stats settings = {} settings["S"] = 1 settings["K"] = k_max settings["channels"] = channels settings["device"] = "cuda" if cuda else "cpu" settings["dtype"] = "double" settings["use_pykeops"] = pykeops # priors settings settings["background_mean_std"] = DEFAULTS["background_mean_std"] settings["background_std_std"] = DEFAULTS["background_std_std"] settings["lamda_rate"] = DEFAULTS["lamda_rate"] settings["height_std"] = DEFAULTS["height_std"] settings["width_min"] = DEFAULTS["width_min"] settings["width_max"] = DEFAULTS["width_max"] settings["proximity_rate"] = DEFAULTS["proximity_rate"] settings["gain_std"] = DEFAULTS["gain_std"] if overwrite: DEFAULTS["cuda"] = cuda DEFAULTS["nbatch-size"] = nbatch_size DEFAULTS["fbatch-size"] = fbatch_size DEFAULTS["learning-rate"] = learning_rate DEFAULTS["matlab"] = matlab with open(cd / ".tapqir" / "config.yaml", "w") as cfg_file: yaml.dump( {key: value for key, value in DEFAULTS.items() if key != "cd"}, cfg_file, sort_keys=False, ) backend = "funsor" if funsor else "pyro" if backend == "pyro": PYRO_BACKEND = "pyro" elif backend == "funsor": import funsor import pyro.contrib.funsor # noqa: F401 funsor.set_backend("torch") PYRO_BACKEND = "contrib.funsor" else: raise ValueError("Only pyro and funsor backends are supported.") with pyro_backend(PYRO_BACKEND): model = models[model](**settings) model.load(cd) model.init(learning_rate, nbatch_size, fbatch_size) typer.echo("Fitting the data ...") exit_code = model.run(num_iter, progress_bar=progress_bar) if exit_code == 0: typer.echo("Fitting the data: Done") elif exit_code == 1: typer.echo("The model hasn't converged!") typer.echo("Computing stats ...") save_stats(model, cd, save_matlab=matlab) typer.echo("Computing stats: Done")