checkpoint = torch.load(CHECKPOINT_PATH) net.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) start_epoch = checkpoint['epoch'] log_string("-> loaded checkpoint %s (epoch: %d)" % (CHECKPOINT_PATH, start_epoch)) # Decay Batchnorm momentum from 0.5 to 0.999 # note: pytorch's BN momentum (default 0.1)= 1 - tensorflow's BN momentum BN_MOMENTUM_INIT = 0.5 BN_MOMENTUM_MAX = 0.001 bn_lbmd = lambda it: max( BN_MOMENTUM_INIT * BN_DECAY_RATE** (int(it / BN_DECAY_STEP)), BN_MOMENTUM_MAX) bnm_scheduler = BNMomentumScheduler(net, bn_lambda=bn_lbmd, last_epoch=start_epoch - 1) def get_current_lr(epoch): lr = BASE_LEARNING_RATE for i, lr_decay_epoch in enumerate(LR_DECAY_STEPS): if epoch >= lr_decay_epoch: lr *= LR_DECAY_RATES[i] return lr def adjust_learning_rate(optimizer, epoch): lr = get_current_lr(epoch) for param_group in optimizer.param_groups: param_group['lr'] = lr
# Load checkpoint if there is any if FLAGS.detector_checkpoint is not None and os.path.isfile(FLAGS.detector_checkpoint): checkpoint = torch.load(FLAGS.detector_checkpoint) pretrained_dict = checkpoint['model_state_dict'] detector.load_state_dict(pretrained_dict) ema_detector.load_state_dict(pretrained_dict) epoch = checkpoint['epoch'] print("Loaded votenet checkpoint %s (epoch: %d)" % (FLAGS.detector_checkpoint, epoch)) # Decay Batchnorm momentum from 0.5 to 0.999 # note: pytorch's BN momentum (default 0.1)= 1 - tensorflow's BN momentum BN_MOMENTUM_INIT = 0.5 BN_MOMENTUM_MAX = 0.001 bn_lbmd = lambda it: max(BN_MOMENTUM_INIT * BN_DECAY_RATE ** (int(it / BN_DECAY_STEP)), BN_MOMENTUM_MAX) bnm_scheduler = BNMomentumScheduler(detector, bn_lambda=bn_lbmd, last_epoch=-1) def get_current_lr(epoch): lr = BASE_LEARNING_RATE for i, lr_decay_epoch in enumerate(LR_DECAY_STEPS): if epoch >= lr_decay_epoch: lr *= LR_DECAY_RATES[i] return lr def adjust_learning_rate(optimizer, epoch): lr = get_current_lr(epoch) for param_group in optimizer.param_groups: param_group['lr'] = lr # TFBoard Visualizers TRAIN_VISUALIZER = TfVisualizer(LOG_DIR, 'train')