def __init__(
self,
task,
model,
dataset,
optimizer,
identifier,
run_dir=None,
is_debug=False,
is_vis=False,
is_hpo=False,
print_every=100,
seed=None,
logger="tensorboard",
local_rank=0,
amp=False,
cpu=False,
name="base_trainer",
):
self.name = name
self.cpu = cpu
self.start_step = 0
if torch.cuda.is_available() and not self.cpu:
self.device = local_rank
else:
self.device = "cpu"
self.cpu = True # handle case when `--cpu` isn't specified
# but there are no gpu devices available
if run_dir is None:
run_dir = os.getcwd()
timestamp = torch.tensor(datetime.datetime.now().timestamp()).to(
self.device)
# create directories from master rank only
distutils.broadcast(timestamp, 0)
timestamp = datetime.datetime.fromtimestamp(
timestamp.int()).strftime("%Y-%m-%d-%H-%M-%S")
if identifier:
timestamp += "-{}".format(identifier)
try:
commit_hash = (subprocess.check_output([
"git",
"-C",
ocpmodels.__path__[0],
"describe",
"--always",
]).strip().decode("ascii"))
# catch instances where code is not being run from a git repo
except Exception:
commit_hash = None
self.config = {
"task": task,
"model": model.pop("name"),
"model_attributes": model,
"optim": optimizer,
"logger": logger,
"amp": amp,
"gpus": distutils.get_world_size() if not self.cpu else 0,
"cmd": {
"identifier": identifier,
"print_every": print_every,
"seed": seed,
"timestamp": timestamp,
"commit": commit_hash,
"checkpoint_dir": os.path.join(run_dir, "checkpoints",
timestamp),
"results_dir": os.path.join(run_dir, "results", timestamp),
"logs_dir": os.path.join(run_dir, "logs", logger, timestamp),
},
}
# AMP Scaler
self.scaler = torch.cuda.amp.GradScaler() if amp else None
if isinstance(dataset, list):
self.config["dataset"] = dataset[0]
if len(dataset) > 1:
self.config["val_dataset"] = dataset[1]
if len(dataset) > 2:
self.config["test_dataset"] = dataset[2]
else:
self.config["dataset"] = dataset
if not is_debug and distutils.is_master() and not is_hpo:
os.makedirs(self.config["cmd"]["checkpoint_dir"], exist_ok=True)
os.makedirs(self.config["cmd"]["results_dir"], exist_ok=True)
os.makedirs(self.config["cmd"]["logs_dir"], exist_ok=True)
self.is_debug = is_debug
self.is_vis = is_vis
self.is_hpo = is_hpo
if self.is_hpo:
# sets the hpo checkpoint frequency
# default is no checkpointing
self.hpo_checkpoint_every = self.config["optim"].get(
"checkpoint_every", -1)
if distutils.is_master():
print(yaml.dump(self.config, default_flow_style=False))
self.load()
self.evaluator = Evaluator(task=name)
def __init__(
self,
task,
model,
dataset,
optimizer,
identifier,
run_dir=None,
is_debug=False,
is_vis=False,
print_every=100,
seed=None,
logger="tensorboard",
local_rank=0,
amp=False,
name="base_trainer",
):
self.name = name
if torch.cuda.is_available():
self.device = local_rank
else:
self.device = "cpu"
if run_dir is None:
run_dir = os.getcwd()
run_dir = Path(run_dir)
timestamp = torch.tensor(datetime.datetime.now().timestamp()).to(
self.device)
# create directories from master rank only
distutils.broadcast(timestamp, 0)
timestamp = datetime.datetime.fromtimestamp(timestamp).strftime(
"%Y-%m-%d-%H-%M-%S")
if identifier:
timestamp += "-{}".format(identifier)
self.config = {
"task": task,
"model": model.pop("name"),
"model_attributes": model,
"optim": optimizer,
"logger": logger,
"amp": amp,
"cmd": {
"identifier": identifier,
"print_every": print_every,
"seed": seed,
"timestamp": timestamp,
"checkpoint_dir": str(run_dir / "checkpoints" / timestamp),
"results_dir": str(run_dir / "results" / timestamp),
"logs_dir": str(run_dir / "logs" / logger / timestamp),
},
}
# AMP Scaler
self.scaler = torch.cuda.amp.GradScaler() if amp else None
if isinstance(dataset, list):
self.config["dataset"] = dataset[0]
if len(dataset) > 1:
self.config["val_dataset"] = dataset[1]
if len(dataset) > 2:
self.config["test_dataset"] = dataset[2]
else:
self.config["dataset"] = dataset
if not is_debug and distutils.is_master():
os.makedirs(self.config["cmd"]["checkpoint_dir"], exist_ok=True)
os.makedirs(self.config["cmd"]["results_dir"], exist_ok=True)
os.makedirs(self.config["cmd"]["logs_dir"], exist_ok=True)
self.is_debug = is_debug
self.is_vis = is_vis
if distutils.is_master():
print(yaml.dump(self.config, default_flow_style=False))
self.load()
self.evaluator = Evaluator(task=name)
def __init__(
self,
task,
model,
dataset,
optimizer,
identifier,
normalizer=None,
timestamp_id=None,
run_dir=None,
is_debug=False,
is_vis=False,
is_hpo=False,
print_every=100,
seed=None,
logger="tensorboard",
local_rank=0,
amp=False,
cpu=False,
name="base_trainer",
slurm={},
):
self.name = name
self.cpu = cpu
self.epoch = 0
self.step = 0
if torch.cuda.is_available() and not self.cpu:
self.device = torch.device(f"cuda:{local_rank}")
else:
self.device = torch.device("cpu")
self.cpu = True # handle case when `--cpu` isn't specified
# but there are no gpu devices available
if run_dir is None:
run_dir = os.getcwd()
if timestamp_id is None:
timestamp = torch.tensor(datetime.datetime.now().timestamp()).to(
self.device
)
# create directories from master rank only
distutils.broadcast(timestamp, 0)
timestamp = datetime.datetime.fromtimestamp(
timestamp.int()
).strftime("%Y-%m-%d-%H-%M-%S")
if identifier:
self.timestamp_id = f"{timestamp}-{identifier}"
else:
self.timestamp_id = timestamp
else:
self.timestamp_id = timestamp_id
try:
commit_hash = (
subprocess.check_output(
[
"git",
"-C",
ocpmodels.__path__[0],
"describe",
"--always",
]
)
.strip()
.decode("ascii")
)
# catch instances where code is not being run from a git repo
except Exception:
commit_hash = None
self.config = {
"task": task,
"model": model.pop("name"),
"model_attributes": model,
"optim": optimizer,
"logger": logger,
"amp": amp,
"gpus": distutils.get_world_size() if not self.cpu else 0,
"cmd": {
"identifier": identifier,
"print_every": print_every,
"seed": seed,
"timestamp_id": self.timestamp_id,
"commit": commit_hash,
"checkpoint_dir": os.path.join(
run_dir, "checkpoints", self.timestamp_id
),
"results_dir": os.path.join(
run_dir, "results", self.timestamp_id
),
"logs_dir": os.path.join(
run_dir, "logs", logger, self.timestamp_id
),
},
"slurm": slurm,
}
# AMP Scaler
self.scaler = torch.cuda.amp.GradScaler() if amp else None
if "SLURM_JOB_ID" in os.environ and "folder" in self.config["slurm"]:
self.config["slurm"]["job_id"] = os.environ["SLURM_JOB_ID"]
self.config["slurm"]["folder"] = self.config["slurm"][
"folder"
].replace("%j", self.config["slurm"]["job_id"])
if isinstance(dataset, list):
if len(dataset) > 0:
self.config["dataset"] = dataset[0]
if len(dataset) > 1:
self.config["val_dataset"] = dataset[1]
if len(dataset) > 2:
self.config["test_dataset"] = dataset[2]
elif isinstance(dataset, dict):
self.config["dataset"] = dataset.get("train", None)
self.config["val_dataset"] = dataset.get("val", None)
self.config["test_dataset"] = dataset.get("test", None)
else:
self.config["dataset"] = dataset
self.normalizer = normalizer
# This supports the legacy way of providing norm parameters in dataset
if self.config.get("dataset", None) is not None and normalizer is None:
self.normalizer = self.config["dataset"]
if not is_debug and distutils.is_master() and not is_hpo:
os.makedirs(self.config["cmd"]["checkpoint_dir"], exist_ok=True)
os.makedirs(self.config["cmd"]["results_dir"], exist_ok=True)
os.makedirs(self.config["cmd"]["logs_dir"], exist_ok=True)
self.is_debug = is_debug
self.is_vis = is_vis
self.is_hpo = is_hpo
if self.is_hpo:
# conditional import is necessary for checkpointing
from ray import tune
from ocpmodels.common.hpo_utils import tune_reporter
# sets the hpo checkpoint frequency
# default is no checkpointing
self.hpo_checkpoint_every = self.config["optim"].get(
"checkpoint_every", -1
)
if distutils.is_master():
print(yaml.dump(self.config, default_flow_style=False))
self.load()
self.evaluator = Evaluator(task=name)
