code: def main(args):
      pyro.set_rng_seed(args.rng_seed)
  
      # 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
  
      # Construct model.
      latent_seq_length = args.latent_seq_length
      if latent_seq_length is None:
          latent_seq_length = int(dataset.max_length * 1.1)
      model = ProfileHMM(
          latent_seq_length,
          dataset.alphabet_length,
          prior_scale=args.prior_scale,
          indel_prior_bias=args.indel_prior_bias,
          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, n_epochs, 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))
  
      # Plots.
      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,
                      "ProfileHMM_plot.loss_{}.pdf".format(time_stamp),
                  )
              )
  
          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,
                      "ProfileHMM_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,
                      "ProfileHMM_plot.delete_prob_{}.pdf".format(
                          time_stamp
                      ),
                  )
              )
  
      if not args.no_save:
          pyro.get_param_store().save(
              os.path.join(
                  args.out_folder,
                  "ProfileHMM_results.params_{}.out".format(time_stamp),
              )
          )
          with open(
              os.path.join(
                  args.out_folder,
                  "ProfileHMM_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
                  )
              )
          with open(
              os.path.join(
                  args.out_folder,
                  "ProfileHMM_results.input_{}.txt".format(time_stamp),
              ),
              "w",
          ) as ow:
              ow.write("[args]\n")
              args.latent_seq_length = model.latent_seq_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/ProfileHMM.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 ProfileHMM
  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