-
Notifications
You must be signed in to change notification settings - Fork 3
/
train.py
35 lines (24 loc) · 822 Bytes
/
train.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
import os
import torch
from dkt import trainer
from dkt.utils import setSeeds
from dkt.dataloader import Preprocess
from args import parse_args
# import wandb
def main(args):
# wandb.login()
setSeeds(42)
device = "cuda" if torch.cuda.is_available() else "cpu"
args.device = device
print(f"MODEL : {args.config} - {args.device}")
preprocess = Preprocess(args)
preprocess.load_train_data(args.file_name)
train_data = preprocess.get_train_data()
train_data, valid_data = preprocess.split_data(train_data)
# wandb.init(project='', entity='', config=vars(args))
# wandb.run.name = args.config
trainer.run(args, train_data, valid_data)
if __name__ == "__main__":
args = parse_args(mode='train')
os.makedirs(args.model_dir, exist_ok=True)
main(args)