iVGAE/utils.py at master · DavidCarlyn/iVGAE · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import random

from argparse import ArgumentParser

import torch
import numpy as np

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

def get_args():
    parser = ArgumentParser()
    parser.add_argument("--seed", type=int, default=2022)
    parser.add_argument("--hidden_channels", type=int, default=10)
    parser.add_argument("--out_channels", type=int, default=10)
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--train_samples", type=int, default=8000)
    parser.add_argument("--val_samples", type=int, default=1000)
    parser.add_argument("--lr", type=float, default=0.001)
    parser.add_argument("--dset", type=str, choices=["asia", "earthquake"], default="asia")
    parser.add_argument("--i_vars", type=str, nargs="+", default=[])
    parser.add_argument("--i_probs", type=float, nargs="+", default=[])

    return parser.parse_args()

def get_test_args():
    parser = ArgumentParser()
    parser.add_argument("--seed", type=int, default=2022)
    parser.add_argument("--hidden_channels", type=int, default=10)
    parser.add_argument("--out_channels", type=int, default=10)
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--test_samples", type=int, default=1000)
    parser.add_argument("--num_of_runs", type=int, default=50)
    parser.add_argument("--dset", type=str, choices=["asia", "earthquake"], default="asia")
    parser.add_argument("--model", type=str, default=None)
    parser.add_argument("--i_vars", type=str, nargs="+", default=[])
    parser.add_argument("--i_probs", type=float, nargs="+", default=[])

    return parser.parse_args()