"""Stage-1: pretrain LED denoiser on SDD.""" import argparse from trainer.train_sdd_pretrain import Trainer def main(): p = argparse.ArgumentParser() p.add_argument('--cuda', default=True) p.add_argument('--cfg', default='sdd/sdd') p.add_argument('--gpu', type=int, default=0) p.add_argument('--info', type=str, default='pretrain') p.add_argument('--grad_accum', type=int, default=32) config = p.parse_args() config.train = 1 Trainer(config).fit() if __name__ == '__main__': main()