| 25 | |
| 26 | |
| 27 | def parse_option(): |
| 28 | parser = argparse.ArgumentParser('argument for training') |
| 29 | |
| 30 | parser.add_argument('--print_freq', type=int, default=10, |
| 31 | help='print frequency') |
| 32 | parser.add_argument('--save_freq', type=int, default=50, |
| 33 | help='save frequency') |
| 34 | parser.add_argument('--batch_size', type=int, default=256, |
| 35 | help='batch_size') |
| 36 | parser.add_argument('--num_workers', type=int, default=16, |
| 37 | help='num of workers to use') |
| 38 | parser.add_argument('--epochs', type=int, default=1000, |
| 39 | help='number of training epochs') |
| 40 | |
| 41 | # optimization |
| 42 | parser.add_argument('--learning_rate', type=float, default=0.05, |
| 43 | help='learning rate') |
| 44 | parser.add_argument('--lr_decay_epochs', type=str, default='700,800,900', |
| 45 | help='where to decay lr, can be a list') |
| 46 | parser.add_argument('--lr_decay_rate', type=float, default=0.1, |
| 47 | help='decay rate for learning rate') |
| 48 | parser.add_argument('--weight_decay', type=float, default=1e-4, |
| 49 | help='weight decay') |
| 50 | parser.add_argument('--momentum', type=float, default=0.9, |
| 51 | help='momentum') |
| 52 | |
| 53 | # model dataset |
| 54 | parser.add_argument('--model', type=str, default='resnet50') |
| 55 | parser.add_argument('--dataset', type=str, default='cifar10', |
| 56 | choices=['cifar10', 'cifar100', 'path'], help='dataset') |
| 57 | parser.add_argument('--mean', type=str, help='mean of dataset in path in form of str tuple') |
| 58 | parser.add_argument('--std', type=str, help='std of dataset in path in form of str tuple') |
| 59 | parser.add_argument('--data_folder', type=str, default=None, help='path to custom dataset') |
| 60 | parser.add_argument('--size', type=int, default=32, help='parameter for RandomResizedCrop') |
| 61 | |
| 62 | # method |
| 63 | parser.add_argument('--method', type=str, default='SupCon', |
| 64 | choices=['SupCon', 'SimCLR'], help='choose method') |
| 65 | |
| 66 | # temperature |
| 67 | parser.add_argument('--temp', type=float, default=0.07, |
| 68 | help='temperature for loss function') |
| 69 | |
| 70 | # other setting |
| 71 | parser.add_argument('--cosine', action='store_true', |
| 72 | help='using cosine annealing') |
| 73 | parser.add_argument('--syncBN', action='store_true', |
| 74 | help='using synchronized batch normalization') |
| 75 | parser.add_argument('--warm', action='store_true', |
| 76 | help='warm-up for large batch training') |
| 77 | parser.add_argument('--trial', type=str, default='0', |
| 78 | help='id for recording multiple runs') |
| 79 | |
| 80 | opt = parser.parse_args() |
| 81 | |
| 82 | # check if dataset is path that passed required arguments |
| 83 | if opt.dataset == 'path': |
| 84 | assert opt.data_folder is not None \ |