(args, model, get_num_layer=None, get_layer_scale=None, filter_bias_and_bn=True, skip_list=None)
| 89 | |
| 90 | |
| 91 | def create_optimizer(args, model, get_num_layer=None, get_layer_scale=None, filter_bias_and_bn=True, skip_list=None): |
| 92 | opt_lower = args.opt.lower() |
| 93 | weight_decay = args.weight_decay |
| 94 | if weight_decay and filter_bias_and_bn: |
| 95 | skip = {} |
| 96 | if skip_list is not None: |
| 97 | skip = skip_list |
| 98 | elif hasattr(model, 'no_weight_decay'): |
| 99 | skip = model.no_weight_decay() |
| 100 | parameters = get_parameter_groups(model, weight_decay, skip, get_num_layer, get_layer_scale) |
| 101 | weight_decay = 0. |
| 102 | else: |
| 103 | parameters = model.parameters() |
| 104 | |
| 105 | if 'fused' in opt_lower: |
| 106 | assert has_apex and torch.cuda.is_available(), 'APEX and CUDA required for fused optimizers' |
| 107 | |
| 108 | opt_args = dict(lr=args.lr, weight_decay=weight_decay) |
| 109 | if hasattr(args, 'opt_eps') and args.opt_eps is not None: |
| 110 | opt_args['eps'] = args.opt_eps |
| 111 | if hasattr(args, 'opt_betas') and args.opt_betas is not None: |
| 112 | opt_args['betas'] = args.opt_betas |
| 113 | |
| 114 | print("optimizer settings:", opt_args) |
| 115 | |
| 116 | opt_split = opt_lower.split('_') |
| 117 | opt_lower = opt_split[-1] |
| 118 | if opt_lower == 'sgd' or opt_lower == 'nesterov': |
| 119 | opt_args.pop('eps', None) |
| 120 | optimizer = optim.SGD(parameters, momentum=args.momentum, nesterov=True, **opt_args) |
| 121 | elif opt_lower == 'momentum': |
| 122 | opt_args.pop('eps', None) |
| 123 | optimizer = optim.SGD(parameters, momentum=args.momentum, nesterov=False, **opt_args) |
| 124 | elif opt_lower == 'adam': |
| 125 | optimizer = optim.Adam(parameters, **opt_args) |
| 126 | elif opt_lower == 'adamw': |
| 127 | optimizer = optim.AdamW(parameters, **opt_args) |
| 128 | elif opt_lower == 'nadam': |
| 129 | optimizer = Nadam(parameters, **opt_args) |
| 130 | elif opt_lower == 'radam': |
| 131 | optimizer = RAdam(parameters, **opt_args) |
| 132 | elif opt_lower == 'adamp': |
| 133 | optimizer = AdamP(parameters, wd_ratio=0.01, nesterov=True, **opt_args) |
| 134 | elif opt_lower == 'sgdp': |
| 135 | optimizer = SGDP(parameters, momentum=args.momentum, nesterov=True, **opt_args) |
| 136 | elif opt_lower == 'adadelta': |
| 137 | optimizer = optim.Adadelta(parameters, **opt_args) |
| 138 | elif opt_lower == 'adafactor': |
| 139 | if not args.lr: |
| 140 | opt_args['lr'] = None |
| 141 | optimizer = Adafactor(parameters, **opt_args) |
| 142 | elif opt_lower == 'adahessian': |
| 143 | optimizer = Adahessian(parameters, **opt_args) |
| 144 | elif opt_lower == 'rmsprop': |
| 145 | optimizer = optim.RMSprop(parameters, alpha=0.9, momentum=args.momentum, **opt_args) |
| 146 | elif opt_lower == 'rmsproptf': |
| 147 | optimizer = RMSpropTF(parameters, alpha=0.9, momentum=args.momentum, **opt_args) |
| 148 | elif opt_lower == 'novograd': |
no test coverage detected