code: def main(args):
      # Load dataset.
      if args.cpu_data or not args.cuda:
          device = torch.device("cpu")
      else:
          device = torch.device("cuda")
      if args.test:
          dataset = generate_data(args.small, args.include_stop, device)
      else:
          dataset = BiosequenceDataset(
              args.file,
              "fasta",
              args.alphabet,
              include_stop=args.include_stop,
              device=device,
          )
      args.batch_size = min([dataset.data_size, args.batch_size])
      if args.split > 0.0:
          # Train test split.
          heldout_num = int(np.ceil(args.split * len(dataset)))
          data_lengths = [len(dataset) - heldout_num, heldout_num]
          # Specific data split seed, for comparability across models and
          # parameter initializations.
          pyro.set_rng_seed(args.rng_data_seed)
          indices = torch.randperm(
              sum(data_lengths), device=device
          ).tolist()
          dataset_train, dataset_test = [
              torch.utils.data.Subset(
                  dataset, indices[(offset - length) : offset]
              )
              for offset, length in zip(
                  torch._utils._accumulate(data_lengths), data_lengths
              )
          ]
      else:
          dataset_train = dataset
          dataset_test = None
  
      # Training seed.
      pyro.set_rng_seed(args.rng_seed)
  
      # Construct model.
      model = FactorMuE(
          dataset.max_length,
          dataset.alphabet_length,
          args.z_dim,
          batch_size=args.batch_size,
          latent_seq_length=args.latent_seq_length,
          indel_factor_dependence=args.indel_factor,
          indel_prior_scale=args.indel_prior_scale,
          indel_prior_bias=args.indel_prior_bias,
          inverse_temp_prior=args.inverse_temp_prior,
          weights_prior_scale=args.weights_prior_scale,
          offset_prior_scale=args.offset_prior_scale,
          z_prior_distribution=args.z_prior,
          ARD_prior=args.ARD_prior,
          substitution_matrix=(not args.no_substitution_matrix),
          substitution_prior_scale=args.substitution_prior_scale,
          latent_alphabet_length=args.latent_alphabet,
          cuda=args.cuda,
          pin_memory=args.pin_mem,
      )
  
      # Infer with SVI.
      scheduler = MultiStepLR(
          {
              "optimizer": Adam,
              "optim_args": {"lr": args.learning_rate},
              "milestones": json.loads(args.milestones),
              "gamma": args.learning_gamma,
          }
      )
      n_epochs = args.n_epochs
      losses = model.fit_svi(
          dataset_train,
          n_epochs,
          args.anneal,
          args.batch_size,
          scheduler,
          args.jit,
      )
  
      # Evaluate.
      train_lp, test_lp, train_perplex, test_perplex = model.evaluate(
          dataset_train, dataset_test, args.jit
      )
      print(
          "train logp: {} perplex: {}".format(train_lp, train_perplex)
      )
      print("test logp: {} perplex: {}".format(test_lp, test_perplex))
  
      # Get latent space embedding.
      z_locs, z_scales = model.embed(dataset)
  
      # Plot and save.
      time_stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
      if not args.no_plots:
          plt.figure(figsize=(6, 6))
          plt.plot(losses)
          plt.xlabel("step")
          plt.ylabel("loss")
          if not args.no_save:
              plt.savefig(
                  os.path.join(
                      args.out_folder,
                      "FactorMuE_plot.loss_{}.pdf".format(time_stamp),
                  )
              )
  
          plt.figure(figsize=(6, 6))
          plt.scatter(z_locs[:, 0], z_locs[:, 1])
          plt.xlabel(r"$z_1$")
          plt.ylabel(r"$z_2$")
          if not args.no_save:
              plt.savefig(
                  os.path.join(
                      args.out_folder,
                      "FactorMuE_plot.latent_{}.pdf".format(time_stamp),
                  )
              )
  
          if not args.indel_factor:
              # Plot indel parameters. See statearrangers.py for details on the
              # r and u parameters.
              plt.figure(figsize=(6, 6))
              insert = pyro.param("insert_q_mn").detach()
              insert_expect = torch.exp(
                  insert - insert.logsumexp(-1, True)
              )
              plt.plot(insert_expect[:, :, 1].cpu().numpy())
              plt.xlabel("position")
              plt.ylabel("probability of insert")
              plt.legend([r"$r_0$", r"$r_1$", r"$r_2$"])
              if not args.no_save:
                  plt.savefig(
                      os.path.join(
                          args.out_folder,
                          "FactorMuE_plot.insert_prob_{}.pdf".format(
                              time_stamp
                          ),
                      )
                  )
              plt.figure(figsize=(6, 6))
              delete = pyro.param("delete_q_mn").detach()
              delete_expect = torch.exp(
                  delete - delete.logsumexp(-1, True)
              )
              plt.plot(delete_expect[:, :, 1].cpu().numpy())
              plt.xlabel("position")
              plt.ylabel("probability of delete")
              plt.legend([r"$u_0$", r"$u_1$", r"$u_2$"])
              if not args.no_save:
                  plt.savefig(
                      os.path.join(
                          args.out_folder,
                          "FactorMuE_plot.delete_prob_{}.pdf".format(
                              time_stamp
                          ),
                      )
                  )
  
      if not args.no_save:
          pyro.get_param_store().save(
              os.path.join(
                  args.out_folder,
                  "FactorMuE_results.params_{}.out".format(time_stamp),
              )
          )
          with open(
              os.path.join(
                  args.out_folder,
                  "FactorMuE_results.evaluation_{}.txt".format(
                      time_stamp
                  ),
              ),
              "w",
          ) as ow:
              ow.write("train_lp,test_lp,train_perplex,test_perplex\n")
              ow.write(
                  "{},{},{},{}\n".format(
                      train_lp, test_lp, train_perplex, test_perplex
                  )
              )
          np.savetxt(
              os.path.join(
                  args.out_folder,
                  "FactorMuE_results.embed_loc_{}.txt".format(
                      time_stamp
                  ),
              ),
              z_locs.cpu().numpy(),
          )
          np.savetxt(
              os.path.join(
                  args.out_folder,
                  "FactorMuE_results.embed_scale_{}.txt".format(
                      time_stamp
                  ),
              ),
              z_scales.cpu().numpy(),
          )
          with open(
              os.path.join(
                  args.out_folder,
                  "FactorMuE_results.input_{}.txt".format(time_stamp),
              ),
              "w",
          ) as ow:
              ow.write("[args]\n")
              args.latent_seq_length = model.latent_seq_length
              args.latent_alphabet = model.latent_alphabet_length
              for elem in list(args.__dict__.keys()):
                  ow.write(
                      "{} = {}\n".format(
                          elem, args.__getattribute__(elem)
                      )
                  )
              ow.write(
                  "alphabet_str = {}\n".format(
                      "".join(dataset.alphabet)
                  )
              )
              ow.write("max_length = {}\n".format(dataset.max_length))
  
file_path: pyro/examples/contrib/mue/FactorMuE.py
imports: import argparse
  import datetime
  import json
  import os
  import matplotlib.pyplot
  import numpy
  import torch
  from torch.optim import Adam
  import pyro
  from pyro.contrib.mue.dataloaders import BiosequenceDataset
  from pyro.contrib.mue.models import FactorMuE
  from pyro.optim import MultiStepLR
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types