device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if args.dataset == 'physionet':
data_obj = utils.get_physionet_data(args, 'cpu', args.quantization)
elif args.dataset == 'mimiciii':
data_obj = utils.get_mimiciii_data(args)
train_loader = data_obj["train_dataloader"]
test_loader = data_obj["test_dataloader"]
val_loader = data_obj["val_dataloader"]
dim = data_obj["input_dim"]
if args.enc == 'enc_rnn3':
rec = models.enc_rnn3(dim,
torch.linspace(0, 1., 128),
args.latent_dim,
args.rec_hidden,
128,
learn_emb=args.learn_emb).to(device)
elif args.enc == 'mtan_rnn':
rec = models.enc_mtan_rnn(dim,
torch.linspace(0, 1., args.num_ref_points),
args.latent_dim,
args.rec_hidden,
embed_time=128,
learn_emb=args.learn_emb,
num_heads=args.enc_num_heads).to(device)
if args.dec == 'rnn3':
dec = models.dec_rnn3(dim,
torch.linspace(0, 1., 128),
args.latent_dim,
if args.dataset == "toy":
data_obj = utils.kernel_smoother_data_gen(args, alpha=100.0, seed=0)
elif args.dataset == "physionet":
data_obj = utils.get_physionet_data(args, "cpu", args.quantization)
train_loader = data_obj["train_dataloader"]
test_loader = data_obj["test_dataloader"]
dim = data_obj["input_dim"]
# model
if args.enc == "enc_rnn3":
rec = models.enc_rnn3(
dim,
torch.linspace(0, 1.0, args.num_ref_points),
args.latent_dim,
args.rec_hidden,
128,
learn_emb=args.learn_emb,
).to(device)
elif args.enc == "mtan_rnn":
rec = models.enc_mtan_rnn(
dim,
torch.linspace(0, 1.0, args.num_ref_points),
args.latent_dim,
args.rec_hidden,
embed_time=128,
learn_emb=args.learn_emb,
num_heads=args.enc_num_heads,
).to(device)
if args.dec == "rnn3":