| 19 | from torchvision.utils import save_image |
| 20 | |
| 21 | def parse_args(): |
| 22 | parser = argparse.ArgumentParser('Model') |
| 23 | parser.add_argument('--batch_size', type=int, default=42, |
| 24 | help='Batch Size during training') |
| 25 | parser.add_argument('--num_workers', type=int, default=0, |
| 26 | help='number of workers for data loader') |
| 27 | parser.add_argument('--epoch', default=30, type=int, help='Epoch to run') |
| 28 | parser.add_argument('--learning_rate', default=0.0001, |
| 29 | type=float, help='Initial learning rate') |
| 30 | parser.add_argument('--log_dir', type=str, default=None, |
| 31 | help='Log path [default: None]') |
| 32 | parser.add_argument('--decay_rate', type=float, |
| 33 | default=1e-4, help='weight decay [default: 1e-4]') |
| 34 | parser.add_argument('--step_size', type=int, default=5, |
| 35 | help='Decay step for lr decay') |
| 36 | parser.add_argument('--lr_decay', type=float, |
| 37 | default=0.5, help='Decay rate for lr decay') |
| 38 | parser.add_argument('--gpu', type=str, default='0,1,2', help='GPU to use') |
| 39 | parser.add_argument('--display', type=int, default=50, |
| 40 | help='number of iteration per display') |
| 41 | parser.add_argument('--env_name', type=str, |
| 42 | default='train_chair', help='enviroment name of visdom') |
| 43 | parser.add_argument('--num_basis', type=int, default=15, |
| 44 | help='number of basis vectors') |
| 45 | return parser.parse_args() |
| 46 | |
| 47 | opt = parse_args() |
| 48 | os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu |