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
|
import logging
import os
import custom_graphgym # noqa, register custom modules
import torch
from torch_geometric import seed_everything
from torch_geometric.graphgym.cmd_args import parse_args
from torch_geometric.graphgym.config import (
cfg,
dump_cfg,
load_cfg,
set_out_dir,
set_run_dir,
)
from torch_geometric.graphgym.logger import set_printing
from torch_geometric.graphgym.model_builder import create_model
from torch_geometric.graphgym.train import GraphGymDataModule, train
from torch_geometric.graphgym.utils.agg_runs import agg_runs
from torch_geometric.graphgym.utils.comp_budget import params_count
from torch_geometric.graphgym.utils.device import auto_select_device
if __name__ == '__main__':
# Load cmd line args
args = parse_args()
# Load config file
load_cfg(cfg, args)
set_out_dir(cfg.out_dir, args.cfg_file)
# Set Pytorch environment
torch.set_num_threads(cfg.num_threads)
dump_cfg(cfg)
# Repeat for different random seeds
for i in range(args.repeat):
set_run_dir(cfg.out_dir)
set_printing()
# Set configurations for each run
cfg.seed = cfg.seed + 1
seed_everything(cfg.seed)
auto_select_device()
# Set machine learning pipeline
datamodule = GraphGymDataModule()
model = create_model()
# Print model info
logging.info(model)
logging.info(cfg)
cfg.params = params_count(model)
logging.info('Num parameters: %s', cfg.params)
train(model, datamodule, logger=True)
# Aggregate results from different seeds
agg_runs(cfg.out_dir, cfg.metric_best)
# When being launched in batch mode, mark a yaml as done
if args.mark_done:
os.rename(args.cfg_file, f'{args.cfg_file}_done')
|
