solver/main_train.py at master · 0x4e8/solver · 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
42
43
44
45
import os
import sys
import torch
import datetime
import argparse
from ml4co.Trainer.MIPtrainer import MIPTrain
from ml4co.DataCollector.MIPdata import MIPData

# sys.path.append(os.path.dirname(os.path.abspath(__file__)))

def main():

    parser = argparse.ArgumentParser(description="A script to handle size and instances parameters")

    parser.add_argument('task', type=str, nargs='?', default='LNS', help='task type (default: branch)')
    parser.add_argument('problem', type=str, nargs='?', default='cauctions', help='Instance type (default: setcover)')
    parser.add_argument('method', type=str, nargs='?', default='gnn', help='Testing methods (default: lns_CL)')

    args = parser.parse_args()

    DATA_DIR =  os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data')
    MODEL_DIR =  os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models')
    DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    print('[green]Starting LNS training...'
          '[/green]')

    print(DATA_DIR, MODEL_DIR, DEVICE)

    MIPTrainer = MIPTrain(model_dir=MODEL_DIR, data_dir=DATA_DIR, device=DEVICE)


    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(args.method, args.problem)
    print(f"[{timestamp}] Starting training | Method: {args.method} | Problem: {args.problem}")

    model_path = MIPTrainer.train(task=args.task, problem=args.problem, method=args.method)
    print(model_path)

    print(f"[{timestamp}] Finished training | Task: {args.task} |Method: {args.method} | Problem: {args.problem}")

if __name__ == '__main__':


    main()