DiffPack/util.py at main · DeepGraphLearning/DiffPack · GitHub

History

141 lines (109 loc) · 4.4 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

import os

import time

import logging

import argparse

import yaml

import jinja2

from jinja2 import meta

import easydict

import torch

from torch import distributed as dist

from torchdrug import core, utils, datasets, models, tasks

from torchdrug.utils import comm

logger = logging.getLogger(__file__)

def get_root_logger(file=True):

logger = logging.getLogger("")

logger.setLevel(logging.INFO)

format = logging.Formatter("%(asctime)-10s %(message)s", "%H:%M:%S")

if file:

handler = logging.FileHandler("log.txt")

handler.setFormatter(format)

logger.addHandler(handler)

return logger

def create_working_directory(cfg):

file_name = "%s_working_dir.tmp" % os.environ["SLURM_JOB_ID"]

world_size = comm.get_world_size()

if world_size > 1 and not dist.is_initialized():

comm.init_process_group("nccl", init_method="env://")

working_dir = os.path.join(os.path.expanduser(cfg.output_dir),

cfg.task["class"], cfg.dataset["class"], cfg.task.model["class"],

time.strftime("%Y-%m-%d-%H-%M-%S"))

# synchronize working directory

if comm.get_rank() == 0:

with open(file_name, "w") as fout:

fout.write(working_dir)

os.makedirs(working_dir)

comm.synchronize()

if comm.get_rank() != 0:

with open(file_name, "r") as fin:

working_dir = fin.read()

comm.synchronize()

if comm.get_rank() == 0:

os.remove(file_name)

os.chdir(working_dir)

return working_dir

def detect_variables(cfg_file):

with open(cfg_file, "r") as fin:

raw = fin.read()

env = jinja2.Environment()

ast = env.parse(raw)

vars = meta.find_undeclared_variables(ast)

return vars

def load_config(cfg_file, context=None):

with open(cfg_file, "r") as fin:

raw = fin.read()

template = jinja2.Template(raw)

instance = template.render(context)

cfg = yaml.safe_load(instance)

cfg = easydict.EasyDict(cfg)

return cfg

def parse_args():

parser = argparse.ArgumentParser()

parser.add_argument("-c", "--config", help="yaml configuration file", required=True)

parser.add_argument("-s", "--seed", help="random seed for PyTorch", type=int, default=1024)

args, unparsed = parser.parse_known_args()

# get dynamic arguments defined in the config file

vars = detect_variables(args.config)

parser = argparse.ArgumentParser()

for var in vars:

parser.add_argument("--%s" % var, default="null")

vars = parser.parse_known_args(unparsed)[0]

vars = {k: utils.literal_eval(v) for k, v in vars._get_kwargs()}

return args, vars

def build_downstream_solver(cfg, dataset):

train_set, valid_set, test_set = dataset.split()

if comm.get_rank() == 0:

logger.warning(dataset)

logger.warning("#train: %d, #valid: %d, #test: %d" % (len(train_set), len(valid_set), len(test_set)))

if cfg.task['class'] == 'MultipleBinaryClassification':

cfg.task.task = [_ for _ in range(len(dataset.tasks))]

else:

cfg.task.task = dataset.tasks

task = core.Configurable.load_config_dict(cfg.task)

if not "lr_ratio" in cfg:

cfg.optimizer.params = task.parameters()

else:

cfg.optimizer.params = [

{'params': task.model.model.parameters(), 'lr': cfg.optimizer.lr * cfg.lr_ratio},

]

cfg.optimizer.params = task.parameters()

optimizer = core.Configurable.load_config_dict(cfg.optimizer)

solver = core.Engine(task, train_set, valid_set, test_set, optimizer, **cfg.engine)

if cfg.get("checkpoint") is not None:

solver.load(cfg.checkpoint)

if cfg.get("model_checkpoint") is not None:

if comm.get_rank() == 0:

logger.warning("Load checkpoint from %s" % cfg.model_checkpoint)

cfg.model_checkpoint = os.path.expanduser(cfg.model_checkpoint)

model_dict = torch.load(cfg.model_checkpoint, map_location=torch.device('cpu'))

task.model.load_state_dict(model_dict)

return solver

def build_pretrain_solver(cfg, dataset):

if comm.get_rank() == 0:

logger.warning(dataset)

logger.warning("#dataset: %d" % (len(dataset)))

task = core.Configurable.load_config_dict(cfg.task)

cfg.optimizer.params = task.parameters()

optimizer = core.Configurable.load_config_dict(cfg.optimizer)

solver = core.Engine(task, dataset, None, None, optimizer, **cfg.engine)

return solver

DiffPack/util.py at main · DeepGraphLearning/DiffPack · GitHub

Uh oh!

Expand file tree

Breadcrumbs

util.py

Latest commit

History

Breadcrumbs

util.py

File metadata and controls