| 22 | |
| 23 | |
| 24 | def parse_option(): |
| 25 | parser = argparse.ArgumentParser('argument for training') |
| 26 | |
| 27 | parser.add_argument('--print_freq', type=int, default=10, |
| 28 | help='print frequency') |
| 29 | parser.add_argument('--save_freq', type=int, default=50, |
| 30 | help='save frequency') |
| 31 | parser.add_argument('--batch_size', type=int, default=256, |
| 32 | help='batch_size') |
| 33 | parser.add_argument('--num_workers', type=int, default=16, |
| 34 | help='num of workers to use') |
| 35 | parser.add_argument('--epochs', type=int, default=100, |
| 36 | help='number of training epochs') |
| 37 | |
| 38 | # optimization |
| 39 | parser.add_argument('--learning_rate', type=float, default=0.1, |
| 40 | help='learning rate') |
| 41 | parser.add_argument('--lr_decay_epochs', type=str, default='60,75,90', |
| 42 | help='where to decay lr, can be a list') |
| 43 | parser.add_argument('--lr_decay_rate', type=float, default=0.2, |
| 44 | help='decay rate for learning rate') |
| 45 | parser.add_argument('--weight_decay', type=float, default=0, |
| 46 | help='weight decay') |
| 47 | parser.add_argument('--momentum', type=float, default=0.9, |
| 48 | help='momentum') |
| 49 | |
| 50 | # model dataset |
| 51 | parser.add_argument('--model', type=str, default='resnet50') |
| 52 | parser.add_argument('--dataset', type=str, default='cifar10', |
| 53 | choices=['cifar10', 'cifar100'], help='dataset') |
| 54 | |
| 55 | # other setting |
| 56 | parser.add_argument('--cosine', action='store_true', |
| 57 | help='using cosine annealing') |
| 58 | parser.add_argument('--warm', action='store_true', |
| 59 | help='warm-up for large batch training') |
| 60 | |
| 61 | parser.add_argument('--ckpt', type=str, default='', |
| 62 | help='path to pre-trained model') |
| 63 | |
| 64 | opt = parser.parse_args() |
| 65 | |
| 66 | # set the path according to the environment |
| 67 | opt.data_folder = './datasets/' |
| 68 | |
| 69 | iterations = opt.lr_decay_epochs.split(',') |
| 70 | opt.lr_decay_epochs = list([]) |
| 71 | for it in iterations: |
| 72 | opt.lr_decay_epochs.append(int(it)) |
| 73 | |
| 74 | opt.model_name = '{}_{}_lr_{}_decay_{}_bsz_{}'.\ |
| 75 | format(opt.dataset, opt.model, opt.learning_rate, opt.weight_decay, |
| 76 | opt.batch_size) |
| 77 | |
| 78 | if opt.cosine: |
| 79 | opt.model_name = '{}_cosine'.format(opt.model_name) |
| 80 | |
| 81 | # warm-up for large-batch training, |