diff --git a/README.md b/README.md index 5ab92987..66cea1d5 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,10 @@ The work of many others is present here. I've tried to make sure all source mate * [Myself](https://github.com/rwightman/pytorch-dpn-pretrained) * LR scheduler ideas from [AllenNLP](https://github.com/allenai/allennlp/tree/master/allennlp/training/learning_rate_schedulers), [FAIRseq](https://github.com/pytorch/fairseq/tree/master/fairseq/optim/lr_scheduler), and SGDR: Stochastic Gradient Descent with Warm Restarts (https://arxiv.org/abs/1608.03983) * Random Erasing from [Zhun Zhong](https://github.com/zhunzhong07/Random-Erasing/blob/master/transforms.py) (https://arxiv.org/abs/1708.04896) - +* Optimizers: + * RAdam by [Liyuan Liu](https://github.com/LiyuanLucasLiu/RAdam) (https://arxiv.org/abs/1908.03265) + * NovoGrad by [Masashi Kimura](https://github.com/convergence-lab/novograd) (https://arxiv.org/abs/1905.11286) + * Lookahead adapted from impl by [Liam](https://github.com/alphadl/lookahead.pytorch) (https://arxiv.org/abs/1907.08610) ## Models I've included a few of my favourite models, but this is not an exhaustive collection. You can't do better than Cadene's collection in that regard. Most models do have pretrained weights from their respective sources or original authors. diff --git a/timm/models/resnet.py b/timm/models/resnet.py index eff83066..89649f7c 100644 --- a/timm/models/resnet.py +++ b/timm/models/resnet.py @@ -44,6 +44,9 @@ default_cfgs = { 'resnet50': _cfg( url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/rw_resnet50-86acaeed.pth', interpolation='bicubic'), + 'resnet50d': _cfg( + url='', + interpolation='bicubic'), 'resnet101': _cfg(url='https://download.pytorch.org/models/resnet101-5d3b4d8f.pth'), 'resnet152': _cfg(url='https://download.pytorch.org/models/resnet152-b121ed2d.pth'), 'tv_resnet34': _cfg(url='https://download.pytorch.org/models/resnet34-333f7ec4.pth'), @@ -259,7 +262,7 @@ class ResNet(nn.Module): def __init__(self, block, layers, num_classes=1000, in_chans=3, use_se=False, cardinality=1, base_width=64, stem_width=64, deep_stem=False, block_reduce_first=1, down_kernel_size=1, avg_down=False, dilated=False, - norm_layer=nn.BatchNorm2d, drop_rate=0.0, global_pool='avg'): + norm_layer=nn.BatchNorm2d, drop_rate=0.0, global_pool='avg', zero_init_last_bn=True): self.num_classes = num_classes self.inplanes = stem_width * 2 if deep_stem else 64 self.cardinality = cardinality @@ -296,11 +299,16 @@ class ResNet(nn.Module): self.num_features = 512 * block.expansion self.fc = nn.Linear(self.num_features * self.global_pool.feat_mult(), num_classes) - for m in self.modules(): + last_bn_name = 'bn3' if 'Bottleneck' in block.__name__ else 'bn2' + for n, m in self.named_modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') elif isinstance(m, nn.BatchNorm2d): - nn.init.constant_(m.weight, 1.) + if zero_init_last_bn and 'layer' in n and last_bn_name in n: + # Initialize weight/gamma of last BN in each residual block to zero + nn.init.constant_(m.weight, 0.) + else: + nn.init.constant_(m.weight, 1.) nn.init.constant_(m.bias, 0.) def _make_layer(self, block, planes, blocks, stride=1, dilation=1, reduce_first=1, @@ -434,6 +442,20 @@ def resnet50(pretrained=False, num_classes=1000, in_chans=3, **kwargs): return model +@register_model +def resnet50d(pretrained=False, num_classes=1000, in_chans=3, **kwargs): + """Constructs a ResNet-50-D model. + """ + default_cfg = default_cfgs['resnet50d'] + model = ResNet( + Bottleneck, [3, 4, 6, 3], stem_width=32, deep_stem=True, avg_down=True, + num_classes=num_classes, in_chans=in_chans, **kwargs) + model.default_cfg = default_cfg + if pretrained: + load_pretrained(model, default_cfg, num_classes, in_chans) + return model + + @register_model def resnet101(pretrained=False, num_classes=1000, in_chans=3, **kwargs): """Constructs a ResNet-101 model. diff --git a/timm/optim/__init__.py b/timm/optim/__init__.py index fe97041d..3213cd68 100644 --- a/timm/optim/__init__.py +++ b/timm/optim/__init__.py @@ -1,3 +1,7 @@ from .nadam import Nadam from .rmsprop_tf import RMSpropTF +from .adamw import AdamW +from .radam import RAdam +from .novograd import NovoGrad +from .lookahead import Lookahead from .optim_factory import create_optimizer diff --git a/timm/optim/adamw.py b/timm/optim/adamw.py new file mode 100644 index 00000000..66f9a959 --- /dev/null +++ b/timm/optim/adamw.py @@ -0,0 +1,117 @@ +""" AdamW Optimizer +Impl copied from PyTorch master +""" +import math +import torch +from torch.optim.optimizer import Optimizer + + +class AdamW(Optimizer): + r"""Implements AdamW algorithm. + + The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_. + The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_. + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups + lr (float, optional): learning rate (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay (float, optional): weight decay coefficient (default: 1e-2) + amsgrad (boolean, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) + + .. _Adam\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _Decoupled Weight Decay Regularization: + https://arxiv.org/abs/1711.05101 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + """ + + def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, + weight_decay=1e-2, amsgrad=False): + if not 0.0 <= lr: + raise ValueError("Invalid learning rate: {}".format(lr)) + if not 0.0 <= eps: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= betas[0] < 1.0: + raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) + if not 0.0 <= betas[1] < 1.0: + raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) + defaults = dict(lr=lr, betas=betas, eps=eps, + weight_decay=weight_decay, amsgrad=amsgrad) + super(AdamW, self).__init__(params, defaults) + + def __setstate__(self, state): + super(AdamW, self).__setstate__(state) + for group in self.param_groups: + group.setdefault('amsgrad', False) + + def step(self, closure=None): + """Performs a single optimization step. + + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group['params']: + if p.grad is None: + continue + + # Perform stepweight decay + p.data.mul_(1 - group['lr'] * group['weight_decay']) + + # Perform optimization step + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead') + amsgrad = group['amsgrad'] + + state = self.state[p] + + # State initialization + if len(state) == 0: + state['step'] = 0 + # Exponential moving average of gradient values + state['exp_avg'] = torch.zeros_like(p.data) + # Exponential moving average of squared gradient values + state['exp_avg_sq'] = torch.zeros_like(p.data) + if amsgrad: + # Maintains max of all exp. moving avg. of sq. grad. values + state['max_exp_avg_sq'] = torch.zeros_like(p.data) + + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + if amsgrad: + max_exp_avg_sq = state['max_exp_avg_sq'] + beta1, beta2 = group['betas'] + + state['step'] += 1 + bias_correction1 = 1 - beta1 ** state['step'] + bias_correction2 = 1 - beta2 ** state['step'] + + # Decay the first and second moment running average coefficient + exp_avg.mul_(beta1).add_(1 - beta1, grad) + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + if amsgrad: + # Maintains the maximum of all 2nd moment running avg. till now + torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq) + # Use the max. for normalizing running avg. of gradient + denom = (max_exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps']) + else: + denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps']) + + step_size = group['lr'] / bias_correction1 + + p.data.addcdiv_(-step_size, exp_avg, denom) + + return loss diff --git a/timm/optim/lookahead.py b/timm/optim/lookahead.py new file mode 100644 index 00000000..cc1fb495 --- /dev/null +++ b/timm/optim/lookahead.py @@ -0,0 +1,88 @@ +""" Lookahead Optimizer Wrapper. +Implementation modified from: https://github.com/alphadl/lookahead.pytorch +Paper: `Lookahead Optimizer: k steps forward, 1 step back` - https://arxiv.org/abs/1907.08610 +""" +import torch +from torch.optim.optimizer import Optimizer +from collections import defaultdict + + +class Lookahead(Optimizer): + def __init__(self, base_optimizer, alpha=0.5, k=6): + if not 0.0 <= alpha <= 1.0: + raise ValueError(f'Invalid slow update rate: {alpha}') + if not 1 <= k: + raise ValueError(f'Invalid lookahead steps: {k}') + self.alpha = alpha + self.k = k + self.base_optimizer = base_optimizer + self.param_groups = self.base_optimizer.param_groups + self.defaults = base_optimizer.defaults + self.state = defaultdict(dict) + for group in self.param_groups: + group["step_counter"] = 0 + + def update_slow_weights(self, group): + for fast_p in group["params"]: + if fast_p.grad is None: + continue + param_state = self.state[fast_p] + if "slow_buffer" not in param_state: + param_state["slow_buffer"] = torch.empty_like(fast_p.data) + param_state["slow_buffer"].copy_(fast_p.data) + slow = param_state["slow_buffer"] + slow.add_(self.alpha, fast_p.data - slow) + fast_p.data.copy_(slow) + + def sync_lookahead(self): + for group in self.param_groups: + self.update_slow_weights(group) + + def step(self, closure=None): + loss = self.base_optimizer.step(closure) + for group in self.param_groups: + group['step_counter'] += 1 + if group['step_counter'] % self.k == 0: + self.update_slow_weights(group) + return loss + + def state_dict(self): + fast_state_dict = self.base_optimizer.state_dict() + slow_state = { + (id(k) if isinstance(k, torch.Tensor) else k): v + for k, v in self.state.items() + } + fast_state = fast_state_dict["state"] + param_groups = fast_state_dict["param_groups"] + return { + "state": fast_state, + "slow_state": slow_state, + "param_groups": param_groups, + } + + def load_state_dict(self, state_dict): + if 'slow_state' not in state_dict: + print('Loading state_dict from optimizer without Lookahead applied') + state_dict['slow_state'] = defaultdict(dict) + slow_state_dict = { + "state": state_dict["slow_state"], + "param_groups": state_dict["param_groups"], + } + fast_state_dict = { + "state": state_dict["state"], + "param_groups": state_dict["param_groups"], + } + super(Lookahead, self).load_state_dict(slow_state_dict) + self.base_optimizer.load_state_dict(fast_state_dict) + + def add_param_group(self, param_group): + r"""Add a param group to the :class:`Optimizer` s `param_groups`. + This can be useful when fine tuning a pre-trained network as frozen + layers can be made trainable and added to the :class:`Optimizer` as + training progresses. + Args: + param_group (dict): Specifies what Tensors should be optimized along + with group specific optimization options. + """ + param_group['step_counter'] = 0 + self.base_optimizer.add_param_group(param_group) diff --git a/timm/optim/novograd.py b/timm/optim/novograd.py new file mode 100644 index 00000000..4137c6aa --- /dev/null +++ b/timm/optim/novograd.py @@ -0,0 +1,77 @@ +"""NovoGrad Optimizer. +Original impl by Masashi Kimura (Convergence Lab): https://github.com/convergence-lab/novograd +Paper: `Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks` + - https://arxiv.org/abs/1905.11286 +""" + +import torch +from torch.optim.optimizer import Optimizer +import math + + +class NovoGrad(Optimizer): + def __init__(self, params, grad_averaging=False, lr=0.1, betas=(0.95, 0.98), eps=1e-8, weight_decay=0): + defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) + super(NovoGrad, self).__init__(params, defaults) + self._lr = lr + self._beta1 = betas[0] + self._beta2 = betas[1] + self._eps = eps + self._wd = weight_decay + self._grad_averaging = grad_averaging + + self._momentum_initialized = False + + def step(self, closure=None): + loss = None + if closure is not None: + loss = closure() + + if not self._momentum_initialized: + for group in self.param_groups: + for p in group['params']: + if p.grad is None: + continue + state = self.state[p] + grad = p.grad.data + if grad.is_sparse: + raise RuntimeError('NovoGrad does not support sparse gradients') + + v = torch.norm(grad)**2 + m = grad/(torch.sqrt(v) + self._eps) + self._wd * p.data + state['step'] = 0 + state['v'] = v + state['m'] = m + state['grad_ema'] = None + self._momentum_initialized = True + + for group in self.param_groups: + for p in group['params']: + if p.grad is None: + continue + state = self.state[p] + state['step'] += 1 + + step, v, m = state['step'], state['v'], state['m'] + grad_ema = state['grad_ema'] + + grad = p.grad.data + g2 = torch.norm(grad)**2 + grad_ema = g2 if grad_ema is None else grad_ema * \ + self._beta2 + g2 * (1. - self._beta2) + grad *= 1.0 / (torch.sqrt(grad_ema) + self._eps) + + if self._grad_averaging: + grad *= (1. - self._beta1) + + g2 = torch.norm(grad)**2 + v = self._beta2*v + (1. - self._beta2)*g2 + m = self._beta1*m + (grad / (torch.sqrt(v) + self._eps) + self._wd * p.data) + bias_correction1 = 1 - self._beta1 ** step + bias_correction2 = 1 - self._beta2 ** step + step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1 + + state['v'], state['m'] = v, m + state['grad_ema'] = grad_ema + p.data.add_(-step_size, m) + return loss diff --git a/timm/optim/optim_factory.py b/timm/optim/optim_factory.py index 7fe3e1e4..c51bdf20 100644 --- a/timm/optim/optim_factory.py +++ b/timm/optim/optim_factory.py @@ -1,5 +1,5 @@ from torch import optim as optim -from timm.optim import Nadam, RMSpropTF +from timm.optim import Nadam, RMSpropTF, AdamW, RAdam, NovoGrad, Lookahead def add_weight_decay(model, weight_decay=1e-5, skip_list=()): @@ -18,35 +18,55 @@ def add_weight_decay(model, weight_decay=1e-5, skip_list=()): def create_optimizer(args, model, filter_bias_and_bn=True): + opt_lower = args.opt.lower() weight_decay = args.weight_decay + if opt_lower == 'adamw' or opt_lower == 'radam': + # compensate for the way current AdamW and RAdam optimizers + # apply the weight-decay + weight_decay /= args.lr if weight_decay and filter_bias_and_bn: parameters = add_weight_decay(model, weight_decay) weight_decay = 0. else: parameters = model.parameters() - if args.opt.lower() == 'sgd': + opt_split = opt_lower.split('_') + opt_lower = opt_split[-1] + if opt_lower == 'sgd': optimizer = optim.SGD( parameters, lr=args.lr, momentum=args.momentum, weight_decay=weight_decay, nesterov=True) - elif args.opt.lower() == 'adam': + elif opt_lower == 'adam': optimizer = optim.Adam( parameters, lr=args.lr, weight_decay=weight_decay, eps=args.opt_eps) - elif args.opt.lower() == 'nadam': + elif opt_lower == 'adamw': + optimizer = AdamW( + parameters, lr=args.lr, weight_decay=weight_decay, eps=args.opt_eps) + elif opt_lower == 'nadam': optimizer = Nadam( parameters, lr=args.lr, weight_decay=weight_decay, eps=args.opt_eps) - elif args.opt.lower() == 'adadelta': + elif opt_lower == 'radam': + optimizer = RAdam( + parameters, lr=args.lr, weight_decay=weight_decay, eps=args.opt_eps) + elif opt_lower == 'adadelta': optimizer = optim.Adadelta( parameters, lr=args.lr, weight_decay=weight_decay, eps=args.opt_eps) - elif args.opt.lower() == 'rmsprop': + elif opt_lower == 'rmsprop': optimizer = optim.RMSprop( parameters, lr=args.lr, alpha=0.9, eps=args.opt_eps, momentum=args.momentum, weight_decay=weight_decay) - elif args.opt.lower() == 'rmsproptf': + elif opt_lower == 'rmsproptf': optimizer = RMSpropTF( parameters, lr=args.lr, alpha=0.9, eps=args.opt_eps, momentum=args.momentum, weight_decay=weight_decay) + elif opt_lower == 'novograd': + optimizer = NovoGrad(parameters, lr=args.lr, weight_decay=weight_decay, eps=args.opt_eps) else: assert False and "Invalid optimizer" raise ValueError + + if len(opt_split) > 1: + if opt_split[0] == 'lookahead': + optimizer = Lookahead(optimizer) + return optimizer diff --git a/timm/optim/radam.py b/timm/optim/radam.py new file mode 100644 index 00000000..9987a334 --- /dev/null +++ b/timm/optim/radam.py @@ -0,0 +1,152 @@ +"""RAdam Optimizer. +Implementation lifted from: https://github.com/LiyuanLucasLiu/RAdam +Paper: `On the Variance of the Adaptive Learning Rate and Beyond` - https://arxiv.org/abs/1908.03265 +""" +import math +import torch +from torch.optim.optimizer import Optimizer, required + + +class RAdam(Optimizer): + + def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0): + defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) + self.buffer = [[None, None, None] for ind in range(10)] + super(RAdam, self).__init__(params, defaults) + + def __setstate__(self, state): + super(RAdam, self).__setstate__(state) + + def step(self, closure=None): + + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + + for p in group['params']: + if p.grad is None: + continue + grad = p.grad.data.float() + if grad.is_sparse: + raise RuntimeError('RAdam does not support sparse gradients') + + p_data_fp32 = p.data.float() + + state = self.state[p] + + if len(state) == 0: + state['step'] = 0 + state['exp_avg'] = torch.zeros_like(p_data_fp32) + state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) + else: + state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) + state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) + + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + beta1, beta2 = group['betas'] + + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + exp_avg.mul_(beta1).add_(1 - beta1, grad) + + state['step'] += 1 + buffered = self.buffer[int(state['step'] % 10)] + if state['step'] == buffered[0]: + N_sma, step_size = buffered[1], buffered[2] + else: + buffered[0] = state['step'] + beta2_t = beta2 ** state['step'] + N_sma_max = 2 / (1 - beta2) - 1 + N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t) + buffered[1] = N_sma + + # more conservative since it's an approximated value + if N_sma >= 5: + step_size = group['lr'] * math.sqrt( + (1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / ( + N_sma_max - 2)) / (1 - beta1 ** state['step']) + else: + step_size = group['lr'] / (1 - beta1 ** state['step']) + buffered[2] = step_size + + if group['weight_decay'] != 0: + p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) + + # more conservative since it's an approximated value + if N_sma >= 5: + denom = exp_avg_sq.sqrt().add_(group['eps']) + p_data_fp32.addcdiv_(-step_size, exp_avg, denom) + else: + p_data_fp32.add_(-step_size, exp_avg) + + p.data.copy_(p_data_fp32) + + return loss + + +class PlainRAdam(Optimizer): + + def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0): + defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) + + super(PlainRAdam, self).__init__(params, defaults) + + def __setstate__(self, state): + super(PlainRAdam, self).__setstate__(state) + + def step(self, closure=None): + + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + + for p in group['params']: + if p.grad is None: + continue + grad = p.grad.data.float() + if grad.is_sparse: + raise RuntimeError('RAdam does not support sparse gradients') + + p_data_fp32 = p.data.float() + + state = self.state[p] + + if len(state) == 0: + state['step'] = 0 + state['exp_avg'] = torch.zeros_like(p_data_fp32) + state['exp_avg_sq'] = torch.zeros_like(p_data_fp32) + else: + state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32) + state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32) + + exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] + beta1, beta2 = group['betas'] + + exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) + exp_avg.mul_(beta1).add_(1 - beta1, grad) + + state['step'] += 1 + beta2_t = beta2 ** state['step'] + N_sma_max = 2 / (1 - beta2) - 1 + N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t) + + if group['weight_decay'] != 0: + p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32) + + # more conservative since it's an approximated value + if N_sma >= 5: + step_size = group['lr'] * math.sqrt( + (1 - beta2_t) * (N_sma - 4) / (N_sma_max - 4) * (N_sma - 2) / N_sma * N_sma_max / ( + N_sma_max - 2)) / (1 - beta1 ** state['step']) + denom = exp_avg_sq.sqrt().add_(group['eps']) + p_data_fp32.addcdiv_(-step_size, exp_avg, denom) + else: + step_size = group['lr'] / (1 - beta1 ** state['step']) + p_data_fp32.add_(-step_size, exp_avg) + + p.data.copy_(p_data_fp32) + + return loss diff --git a/timm/scheduler/scheduler_factory.py b/timm/scheduler/scheduler_factory.py index 8f1032a1..80d37b37 100644 --- a/timm/scheduler/scheduler_factory.py +++ b/timm/scheduler/scheduler_factory.py @@ -5,6 +5,7 @@ from .step_lr import StepLRScheduler def create_scheduler(args, optimizer): num_epochs = args.epochs + lr_scheduler = None #FIXME expose cycle parms of the scheduler config to arguments if args.sched == 'cosine': lr_scheduler = CosineLRScheduler( @@ -31,7 +32,7 @@ def create_scheduler(args, optimizer): t_in_epochs=True, ) num_epochs = lr_scheduler.get_cycle_length() + args.cooldown_epochs - else: + elif args.sched == 'step': lr_scheduler = StepLRScheduler( optimizer, decay_t=args.decay_epochs, diff --git a/train.py b/train.py index c795927d..9c75e050 100644 --- a/train.py +++ b/train.py @@ -251,7 +251,7 @@ def main(): start_epoch = args.start_epoch elif resume_epoch is not None: start_epoch = resume_epoch - if start_epoch > 0: + if lr_scheduler is not None and start_epoch > 0: lr_scheduler.step(start_epoch) if args.local_rank == 0: @@ -285,10 +285,12 @@ def main(): collate_fn=collate_fn, ) - eval_dir = os.path.join(args.data, 'validation') + eval_dir = os.path.join(args.data, 'val') if not os.path.isdir(eval_dir): - logging.error('Validation folder does not exist at: {}'.format(eval_dir)) - exit(1) + eval_dir = os.path.join(args.data, 'validation') + if not os.path.isdir(eval_dir): + logging.error('Validation folder does not exist at: {}'.format(eval_dir)) + exit(1) dataset_eval = Dataset(eval_dir) loader_eval = create_loader( @@ -390,8 +392,7 @@ def train_epoch( last_batch = batch_idx == last_idx data_time_m.update(time.time() - end) if not args.prefetcher: - input = input.cuda() - target = target.cuda() + input, target = input.cuda(), target.cuda() if args.mixup > 0.: lam = 1. if not args.mixup_off_epoch or epoch < args.mixup_off_epoch: @@ -461,6 +462,10 @@ def train_epoch( lr_scheduler.step_update(num_updates=num_updates, metric=losses_m.avg) end = time.time() + # end for + + if hasattr(optimizer, 'sync_lookahead'): + optimizer.sync_lookahead() return OrderedDict([('loss', losses_m.avg)])