File size: 1,822 Bytes
352cafd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
from argparse import ArgumentParser
class HyperParameters():
def parse(self, unknown_arg_ok=False):
parser = ArgumentParser()
# Generic learning parameters
parser.add_argument('-i', '--iterations', help='Number of training iterations', default=4.5e4, type=int)
parser.add_argument('-b', '--batch_size', help='Batch size', default=12, type=int)
parser.add_argument('--lr', help='Initial learning rate', default=2.25e-4, type=float)
parser.add_argument('--steps', help='Iteration at which learning rate is decayed by gamma', default=[22500, 37500], type=int, nargs='*')
parser.add_argument('--gamma', help='Gamma used in learning rate decay', default=0.1, type=float)
parser.add_argument('--weight_decay', help='Weight decay', default=1e-4, type=float)
# same decay applied to discriminator
parser.add_argument('--load', help='Path to pretrained model if available')
parser.add_argument('--ce_weight', help='Weight of the CE loss', default=0.0, type=float)
parser.add_argument('--l1_weight', help='Weight of the L1 loss', default=1.0, type=float)
parser.add_argument('--l2_weight', help='Weight of the L2 loss', default=1.0, type=float)
parser.add_argument('--grad_weight', help='Weight of the gradient loss', default=5.0, type=float)
# Logging information, this one is positional and mandatory
parser.add_argument('id', help='Experiment UNIQUE id, use NULL to disable logging to tensorboard')
if unknown_arg_ok:
args, _ = parser.parse_known_args()
self.args = vars(args)
else:
self.args = vars(parser.parse_args())
def __getitem__(self, key):
return self.args[key]
def __str__(self):
return str(self.args)
|