def noamwd_decay(step, warmup_steps,
                 model_size, rate=0.5, decay_steps=1000, start_step=500):
    """Learning rate schedule optimized for huge batches
    """
    return (
        model_size ** (-0.5) *
        min(step ** (-0.5), step * warmup_steps**(-1.5)) *
        rate ** (max(step - start_step + decay_steps, 0) // decay_steps))
Learning rate schedule optimized for huge batches
    
lsp_model/optim.py