File size: 808 Bytes
d4cbafd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | """Stage-2: train LED initializer on SDD (baseline or +graph)."""
import argparse
from trainer.train_sdd_led 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='baseline')
p.add_argument('--learning_rate', type=float, default=0.002)
p.add_argument('--train', type=int, default=1)
p.add_argument('--grad_accum', type=int, default=16)
p.add_argument('--use_graph', action='store_true')
p.add_argument('--use_v6_graph', action='store_true')
p.add_argument('--residual_on', type=str, default='y0')
config = p.parse_args()
Trainer(config).fit()
if __name__ == '__main__':
main()
|