70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
import argparse
|
|
import datetime
|
|
import importlib
|
|
import os
|
|
from shutil import copyfile
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from omegaconf import OmegaConf
|
|
|
|
from utils.dist_utils import (
|
|
get_world_size,
|
|
)
|
|
from utils.utils import seed_all
|
|
|
|
|
|
parser = argparse.ArgumentParser(description="VFI")
|
|
parser.add_argument("-c", "--config", type=str)
|
|
parser.add_argument("-p", "--port", default="23455", type=str)
|
|
parser.add_argument("--local_rank", default="0")
|
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
def main_worker(rank, config):
|
|
if "local_rank" not in config:
|
|
config["local_rank"] = config["global_rank"] = rank
|
|
if torch.cuda.is_available():
|
|
print(f"Rank {rank} is available")
|
|
config["device"] = f"cuda:{rank}"
|
|
if config["distributed"]:
|
|
dist.init_process_group(backend="nccl", timeout=datetime.timedelta(seconds=5400))
|
|
else:
|
|
config["device"] = "cpu"
|
|
|
|
cfg_name = os.path.basename(args.config).split(".")[0]
|
|
config["exp_name"] = cfg_name + "_" + config["exp_name"]
|
|
config["save_dir"] = os.path.join(config["save_dir"], config["exp_name"])
|
|
|
|
if (not config["distributed"]) or rank == 0:
|
|
os.makedirs(config["save_dir"], exist_ok=True)
|
|
os.makedirs(f"{config['save_dir']}/ckpts", exist_ok=True)
|
|
config_path = os.path.join(config["save_dir"], args.config.split("/")[-1])
|
|
if not os.path.isfile(config_path):
|
|
copyfile(args.config, config_path)
|
|
print("[**] create folder {}".format(config["save_dir"]))
|
|
|
|
trainer_name = config.get("trainer_type", "base_trainer")
|
|
print(f"using GPU {rank} for training")
|
|
if rank == 0:
|
|
print(trainer_name)
|
|
trainer_pack = importlib.import_module("trainers." + trainer_name)
|
|
trainer = trainer_pack.Trainer(config)
|
|
|
|
trainer.train()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
torch.backends.cudnn.benchmark = True
|
|
cfg = OmegaConf.load(args.config)
|
|
seed_all(cfg.seed)
|
|
rank = int(args.local_rank)
|
|
torch.cuda.set_device(torch.device(f"cuda:{rank}"))
|
|
# setting distributed cfgurations
|
|
cfg["world_size"] = get_world_size()
|
|
cfg["local_rank"] = rank
|
|
if rank == 0:
|
|
print("world_size: ", cfg["world_size"])
|
|
main_worker(rank, cfg)
|