LED/main_sdd_pretrain.py

"""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()