File size: 606 Bytes
af7c0ce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
import os.path as osp
import argparse
from omegaconf import OmegaConf
from paintmind.engine.util import instantiate_from_config
from paintmind.utils.device_utils import configure_compute_backend
def train():
configure_compute_backend()
parser = argparse.ArgumentParser()
parser.add_argument('--cfg', type=str, default='configs/vit_vqgan.yaml')
args = parser.parse_args()
cfg_file = args.cfg
assert osp.exists(cfg_file)
config = OmegaConf.load(cfg_file)
trainer = instantiate_from_config(config.trainer)
trainer.train(args.cfg)
if __name__ == '__main__':
train()
|