File size: 524 Bytes
d4cbafd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | """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()
|