Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import math | |
| from .base_options import BaseOptions | |
| class TrainOptions(BaseOptions): | |
| def initialize(self, parser): | |
| parser = BaseOptions.initialize(self, parser) | |
| parser.add_argument('--aug', default='A', type=str, | |
| help='data augmentation for training') | |
| parser.add_argument('--beta', type=float, default=0.5, | |
| help='balance between Jigsaw and InsDis') | |
| parser.add_argument('--warm', action='store_true', | |
| help='add warm-up setting') | |
| parser.add_argument('--amp', action='store_true', | |
| help='using mixed precision') | |
| parser.add_argument('--opt_level', type=str, default='O2', | |
| choices=['O1', 'O2']) | |
| return parser | |
| def modify_options(self, opt): | |
| opt = self.override_options(opt) | |
| iterations = opt.lr_decay_epochs.split(',') | |
| opt.lr_decay_epochs = list([]) | |
| for it in iterations: | |
| opt.lr_decay_epochs.append(int(it)) | |
| # set up saving name | |
| opt.model_name = '{}_{}_{}_Jig_{}_{}_aug_{}_{}_{}'.format( | |
| opt.method, opt.arch, opt.modal, opt.jigsaw, opt.mem, | |
| opt.aug, opt.head, opt.nce_t | |
| ) | |
| if opt.amp: | |
| opt.model_name = '{}_amp_{}'.format(opt.model_name, opt.opt_level) | |
| if opt.cosine: | |
| opt.model_name = '{}_cosine'.format(opt.model_name) | |
| # warm-up for large-batch training, e.g. 1024 with multiple nodes | |
| if opt.batch_size > 256: | |
| opt.warm = True | |
| if opt.warm: | |
| opt.model_name = '{}_warm'.format(opt.model_name) | |
| opt.warmup_from = 0.01 | |
| if opt.epochs > 500: | |
| opt.warm_epochs = 10 | |
| else: | |
| opt.warm_epochs = 5 | |
| if opt.cosine: | |
| eta_min = opt.learning_rate * (opt.lr_decay_rate ** 3) | |
| opt.warmup_to = eta_min + (opt.learning_rate - eta_min) * ( | |
| 1 + math.cos(math.pi * opt.warm_epochs / opt.epochs)) / 2 | |
| else: | |
| opt.warmup_to = opt.learning_rate | |
| # create folders | |
| opt.model_folder = os.path.join(opt.model_path, opt.model_name) | |
| if not os.path.isdir(opt.model_folder): | |
| os.makedirs(opt.model_folder) | |
| opt.tb_folder = os.path.join(opt.tb_path, opt.model_name) | |
| if not os.path.isdir(opt.tb_folder): | |
| os.makedirs(opt.tb_folder) | |
| return opt | |