73 lines
2.2 KiB
Python
73 lines
2.2 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
#
|
|
# This source code is licensed under the Apache License, Version 2.0
|
|
# found in the LICENSE file in the root directory of this source tree.
|
|
|
|
import math
|
|
import logging
|
|
import os
|
|
|
|
from omegaconf import OmegaConf
|
|
|
|
import dinov2.distributed as distributed
|
|
from dinov2.logging import setup_logging
|
|
from dinov2.utils import utils
|
|
from dinov2.configs import dinov2_default_config
|
|
|
|
|
|
logger = logging.getLogger("dinov2")
|
|
|
|
|
|
def apply_scaling_rules_to_cfg(cfg): # to fix
|
|
if cfg.optim.scaling_rule == "sqrt_wrt_1024":
|
|
base_lr = cfg.optim.base_lr
|
|
cfg.optim.lr = base_lr
|
|
cfg.optim.lr *= math.sqrt(cfg.train.batch_size_per_gpu * distributed.get_global_size() / 1024.0)
|
|
logger.info(f"sqrt scaling learning rate; base: {base_lr}, new: {cfg.optim.lr}")
|
|
else:
|
|
raise NotImplementedError
|
|
return cfg
|
|
|
|
|
|
def write_config(cfg, output_dir, name="config.yaml"):
|
|
logger.info(OmegaConf.to_yaml(cfg))
|
|
saved_cfg_path = os.path.join(output_dir, name)
|
|
with open(saved_cfg_path, "w") as f:
|
|
OmegaConf.save(config=cfg, f=f)
|
|
return saved_cfg_path
|
|
|
|
|
|
def get_cfg_from_args(args):
|
|
args.output_dir = os.path.abspath(args.output_dir)
|
|
args.opts += [f"train.output_dir={args.output_dir}"]
|
|
default_cfg = OmegaConf.create(dinov2_default_config)
|
|
cfg = OmegaConf.load(args.config_file)
|
|
cfg = OmegaConf.merge(default_cfg, cfg, OmegaConf.from_cli(args.opts))
|
|
return cfg
|
|
|
|
|
|
def default_setup(args):
|
|
distributed.enable(overwrite=True)
|
|
seed = getattr(args, "seed", 0)
|
|
rank = distributed.get_global_rank()
|
|
|
|
global logger
|
|
setup_logging(output=args.output_dir, level=logging.INFO)
|
|
logger = logging.getLogger("dinov2")
|
|
|
|
utils.fix_random_seeds(seed + rank)
|
|
logger.info("git:\n {}\n".format(utils.get_sha()))
|
|
logger.info("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
|
|
|
|
|
|
def setup(args):
|
|
"""
|
|
Create configs and perform basic setups.
|
|
"""
|
|
cfg = get_cfg_from_args(args)
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
default_setup(args)
|
|
apply_scaling_rules_to_cfg(cfg)
|
|
write_config(cfg, args.output_dir)
|
|
return cfg
|