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')