-
Notifications
You must be signed in to change notification settings - Fork 6
/
main.py
83 lines (63 loc) · 2.19 KB
/
main.py
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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import torch, os
import argparse
import numpy as np
import pandas as pd
import json
from datetime import datetime
from easydict import EasyDict as edict
from utils.utils import (
load_config,
setup_seed,
get_trainedNets,
get_test_result,
get_dataloaders,
get_models,
)
setup_seed(42)
def single_run(config, device, log_dir):
result_ls = []
# get data
train_loader, val_loader, test_loader = get_dataloaders(config)
# get modelp
model = get_models(config, device)
# train
model, perf = get_trainedNets(config, model, train_loader, val_loader, device, log_dir)
result_ls.append(perf)
# test
perf, test_df = get_test_result(config, model, test_loader, device)
test_df.to_csv(os.path.join(log_dir, "user_mode_detail.csv"))
result_ls.append(perf)
return result_ls
def init_save_path(config):
"""define the path to save, and save the configuration file."""
if config.networkName == "rnn" and config.attention:
networkName = f"{config.dataset}_{config.networkName}_Attn"
else:
networkName = f"{config.dataset}_{config.networkName}"
log_dir = os.path.join(config.save_root, f"{networkName}_{str(int(datetime.now().timestamp()))}")
if not os.path.exists(log_dir):
os.makedirs(log_dir)
with open(os.path.join(log_dir, "conf.json"), "w") as fp:
json.dump(config, fp, indent=4, sort_keys=True)
return log_dir
if __name__ == "__main__":
# load configs
parser = argparse.ArgumentParser()
parser.add_argument(
"config", type=str, nargs="?", help="Path to the config file.", default="config/geolife/transformer.yml"
)
args = parser.parse_args()
config = load_config(args.config)
config = edict(config)
# save the conf
log_dir = init_save_path(config)
#
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# run
result_ls = []
result_ls.extend(single_run(config, device, log_dir))
# transfer to dataframe and save the results
result_df = pd.DataFrame(result_ls)
# print(result_df)
filename = os.path.join(log_dir, f"{config.dataset}_{config.networkName}.csv")
result_df.to_csv(filename, index=False)