-
Notifications
You must be signed in to change notification settings - Fork 1
/
run.py
71 lines (65 loc) · 2.19 KB
/
run.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
import copy
from pathlib import Path
from datetime import datetime
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
from crystal_gnn.config import ex
from crystal_gnn.datamodules import _datamodules
from crystal_gnn.models import _models
@ex.automain
def main(_config):
_config = copy.deepcopy(_config)
pl.seed_everything(_config["seed"])
project_name = _config["project_name"]
exp_name = _config["exp_name"]
log_dir = Path(_config["log_dir"], _config["source"])
# set datamodule
dm = _datamodules[_config["source"]](_config)
# prepare data
dm.prepare_data()
# set model
_config["mean"] = dm.mean
_config["std"] = dm.std
model = _models[_config["model_name"]](_config)
print(model)
# set checkpoint callback
checkpoint_callback = ModelCheckpoint(
save_top_k=1,
verbose=True,
monitor="val/loss",
mode="min",
filename="best-{epoch}",
)
lr_callback = LearningRateMonitor(logging_interval="step")
callbacks = [checkpoint_callback, lr_callback]
# set logger
logger = WandbLogger(
project=project_name,
name=f"{exp_name}",
version=(
f"{exp_name}_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}"
if not _config["test_only"]
else None
),
save_dir=log_dir,
log_model="True",
group=f"{_config['source']}/{_config['target']}/{_config['model_name']}",
)
# set trainer
trainer = pl.Trainer(
num_nodes=_config["num_nodes"],
devices=_config["devices"],
accelerator=_config["accelerator"],
max_epochs=_config["max_epochs"],
strategy="ddp_find_unused_parameters_true",
deterministic=_config["deterministic"],
callbacks=callbacks,
logger=logger,
)
if not _config["test_only"]:
trainer.fit(model, dm, ckpt_path=_config["resume_from"])
trainer.test(model, datamodule=dm, ckpt_path="best")
else:
print(f"load model from {_config['load_path']}")
trainer.test(model, datamodule=dm, ckpt_path=_config["load_path"])