File size: 1,340 Bytes
46a8d8a |
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 |
import torch.distributed as dist
import torch.multiprocessing as mp
import os
import os.path as osp
import sys
import numpy as np
import copy
from lib.cfg_holder import cfg_unique_holder as cfguh
from lib.cfg_helper import \
get_command_line_args, \
cfg_initiates
from lib.model_zoo.sd import version
from lib.utils import get_obj_from_str
if __name__ == "__main__":
cfg = get_command_line_args()
cfg = cfg_initiates(cfg)
if 'train' in cfg:
trainer = get_obj_from_str(cfg.train.main)(cfg)
tstage = get_obj_from_str(cfg.train.stage)()
if 'eval' in cfg:
tstage.nested_eval_stage = get_obj_from_str(cfg.eval.stage)()
trainer.register_stage(tstage)
if cfg.env.gpu_count == 1:
trainer(0)
else:
mp.spawn(trainer,
args=(),
nprocs=cfg.env.gpu_count,
join=True)
trainer.destroy()
else:
evaler = get_obj_from_str(cfg.eval.main)(cfg)
estage = get_obj_from_str(cfg.eval.stage)()
evaler.register_stage(estage)
if cfg.env.gpu_count == 1:
evaler(0)
else:
mp.spawn(evaler,
args=(),
nprocs=cfg.env.gpu_count,
join=True)
evaler.destroy()
|