learning_rate.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); you
// may not use this file except in compliance with the License.  You may
// obtain a copy of the License at:
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the license.
//
// lbann_learning_rate .hpp .cpp - Callback hooks for learning rate schedules

#ifndef LBANN_CALLBACKS_LEARNING_RATE_HPP_INCLUDED
#define LBANN_CALLBACKS_LEARNING_RATE_HPP_INCLUDED

#include "lbann/callbacks/callback.hpp"
#include <unordered_map>
#include <unordered_set>

namespace lbann {

// Forward declarations
class optimizer;

namespace callback {

// Different schedules should inherit from learning_rate.

class learning_rate : public callback_base
{
public:
  learning_rate();
  learning_rate(const learning_rate&) = default;
  learning_rate& operator=(const learning_rate&) = default;
  learning_rate(std::vector<std::string> weights_names);
  void setup(model* m) override;
  void on_epoch_end(model* m) override;

  using callback_base::on_backward_prop_end;
  void on_backward_prop_end(model* m) override;

protected:
  std::vector<std::string> const& get_weights_names() const
  {
    return m_weights_names;
  }

protected:
  virtual float global_schedule(model* m)
  {
    return get_current_global_learning_rate();
  }

  virtual float optimizer_schedule(model* m, optimizer& opt);

  const std::unordered_set<weights*>& get_weights() const noexcept
  {
    return m_weights;
  }

  static float get_current_global_learning_rate() noexcept
  {
    return m_cur_global_lr;
  }

  static void update_global_learning_rate(float rate) noexcept
  {
    m_cur_global_lr = rate;
  }

private:
  static float m_cur_global_lr;

  std::vector<std::string> m_weights_names;

  std::unordered_set<weights*> m_weights;
};

class step_learning_rate : public learning_rate
{
public:
  step_learning_rate(size_t step, float amt);
  step_learning_rate(size_t step,
                     float amt,
                     std::vector<std::string> weights_names);
  step_learning_rate(const step_learning_rate&) = default;
  step_learning_rate& operator=(const step_learning_rate&) = default;
  step_learning_rate* copy() const override
  {
    return new step_learning_rate(*this);
  }
  std::string name() const override { return "step learning rate"; }

protected:
  float global_schedule(model* m) override;

private:
  void write_specific_proto(lbann_data::Callback& proto) const final;

  size_t m_step;
  float m_amt;
};

// Builder function
std::unique_ptr<callback_base> build_step_learning_rate_callback_from_pbuf(
  const google::protobuf::Message&,
  std::shared_ptr<lbann_summary> const&);

class set_learning_rate : public learning_rate
{
public:
  set_learning_rate(size_t step, float val);
  set_learning_rate(size_t step,
                    float val,
                    std::vector<std::string> weights_names);
  set_learning_rate(const set_learning_rate&) = default;
  set_learning_rate& operator=(const set_learning_rate&) = default;
  set_learning_rate* copy() const override
  {
    return new set_learning_rate(*this);
  }
  std::string name() const override { return "step learning rate"; }

protected:
  float global_schedule(model* m) override;

private:
  void write_specific_proto(lbann_data::Callback& proto) const final;
  size_t m_step;
  float m_val;
};

// Builder function
std::unique_ptr<callback_base> build_set_learning_rate_callback_from_pbuf(
  const google::protobuf::Message&,
  std::shared_ptr<lbann_summary> const&);

class adaptive_learning_rate : public learning_rate
{
public:
  adaptive_learning_rate(size_t patience, float amt);
  adaptive_learning_rate(size_t patience,
                         float amt,
                         std::vector<std::string> weights_names);
  adaptive_learning_rate(const adaptive_learning_rate&) = default;
  adaptive_learning_rate& operator=(const adaptive_learning_rate&) = default;
  adaptive_learning_rate* copy() const override
  {
    return new adaptive_learning_rate(*this);
  }
  std::string name() const override { return "adaptive learning rate"; }

protected:
  float global_schedule(model* m) override;

private:
  void write_specific_proto(lbann_data::Callback& proto) const final;

  size_t m_patience;
  float m_amt;
  size_t m_cur_epoch = std::numeric_limits<size_t>::max();
  EvalType m_last_score = std::numeric_limits<EvalType>::max();
  size_t m_wait = 0;
  bool m_adjust_learning_rate = false;
};

// Builder function
std::unique_ptr<callback_base> build_adaptive_learning_rate_callback_from_pbuf(
  const google::protobuf::Message&,
  std::shared_ptr<lbann_summary> const&);

class drop_fixed_learning_rate : public learning_rate
{
public:
  drop_fixed_learning_rate(std::vector<size_t> drop_epochs, float amt);
  drop_fixed_learning_rate(std::vector<size_t> drop_epochs,
                           float amt,
                           std::vector<std::string> weights_names);
  drop_fixed_learning_rate(const drop_fixed_learning_rate&) = default;
  drop_fixed_learning_rate&
  operator=(const drop_fixed_learning_rate&) = default;
  drop_fixed_learning_rate* copy() const override
  {
    return new drop_fixed_learning_rate(*this);
  }
  std::string name() const override { return "drop fixed learning rate"; }

protected:
  float global_schedule(model* m) override;

private:
  void write_specific_proto(lbann_data::Callback& proto) const final;

  float m_amt;
  std::vector<size_t> m_drop_epochs;
};

// Builder function
std::unique_ptr<callback_base>
build_drop_fixed_learning_rate_callback_from_pbuf(
  const google::protobuf::Message&,
  std::shared_ptr<lbann_summary> const&);

class linear_growth_learning_rate : public learning_rate
{
public:
  linear_growth_learning_rate(float target, size_t num_epochs);
  linear_growth_learning_rate(float target, size_t num_epochs, size_t delay);
  linear_growth_learning_rate(float target,
                              size_t num_epochs,
                              size_t delay,
                              std::vector<std::string> weights_names);
  linear_growth_learning_rate(const linear_growth_learning_rate&) = default;
  linear_growth_learning_rate&
  operator=(const linear_growth_learning_rate&) = default;
  linear_growth_learning_rate* copy() const override
  {
    return new linear_growth_learning_rate(*this);
  }
  void setup(model* m) override;
  std::string name() const override { return "linear growth learning rate"; }

protected:
  float global_schedule(model* m) override;

private:
  void write_specific_proto(lbann_data::Callback& proto) const final;

  float m_base_lr;
  float m_target;
  float m_inc;
  size_t m_num_epochs;
  size_t m_delay;
};

// Builder function
std::unique_ptr<callback_base>
build_linear_growth_learning_rate_callback_from_pbuf(
  const google::protobuf::Message&,
  std::shared_ptr<lbann_summary> const&);

class poly_learning_rate : public learning_rate
{
public:
  poly_learning_rate(double p, size_t n_epochs, size_t max_iter);
  poly_learning_rate(double p,
                     size_t n_epochs,
                     size_t max_iter,
                     double endl_r,
                     std::vector<std::string> weights_names);
  poly_learning_rate(const poly_learning_rate&) = default;
  poly_learning_rate& operator=(const poly_learning_rate&) = default;
  poly_learning_rate* copy() const override
  {
    return new poly_learning_rate(*this);
  }
  void setup(model* m) override;
  std::string name() const override { return "poly learning rate"; }

protected:
  float global_schedule(model* m) override;
  float optimizer_schedule(model* m, optimizer& opt) override;

private:
  void write_specific_proto(lbann_data::Callback& proto) const final;

  double m_p;
  size_t m_num_epochs;
  size_t m_max_iter;
  float m_start_lr;
  float m_end_lr;
};

// Builder function
std::unique_ptr<callback_base> build_poly_learning_rate_callback_from_pbuf(
  const google::protobuf::Message&,
  std::shared_ptr<lbann_summary> const&);

class optimizerwise_adaptive_learning_rate : public learning_rate
{
public:
  optimizerwise_adaptive_learning_rate(float scale);
  optimizerwise_adaptive_learning_rate(float scale,
                                       std::vector<std::string> weights_names);
  optimizerwise_adaptive_learning_rate(
    const optimizerwise_adaptive_learning_rate&) = default;
  optimizerwise_adaptive_learning_rate&
  operator=(const optimizerwise_adaptive_learning_rate&) = default;
  optimizerwise_adaptive_learning_rate* copy() const override
  {
    return new optimizerwise_adaptive_learning_rate(*this);
  }
  std::string name() const override
  {
    return "optimizerwise adaptive learning rate";
  }

protected:
  float optimizer_schedule(model* m, optimizer& opt) override;

private:
  void write_specific_proto(lbann_data::Callback& proto) const final;

  float m_scale;
};

// Builder function
std::unique_ptr<callback_base>
build_optimizerwise_adaptive_learning_rate_callback_from_pbuf(
  const google::protobuf::Message&,
  std::shared_ptr<lbann_summary> const&);

class cosine_decay_learning_rate : public learning_rate
{
public:
  cosine_decay_learning_rate(double lr_max,
                             double lr_min,
                             size_t decay_steps,
                             double initial_learning_rate = 0.0,
                             size_t warmup_steps = 0);
  cosine_decay_learning_rate(double lr_max,
                             double lr_min,
                             size_t decay_steps,
                             double initial_learning_rate,
                             size_t warmup_steps,
                             std::vector<std::string> weight_names);
  cosine_decay_learning_rate(const cosine_decay_learning_rate&) = default;
  cosine_decay_learning_rate&
  operator=(const cosine_decay_learning_rate&) = default;
  cosine_decay_learning_rate* copy() const override
  {
    return new cosine_decay_learning_rate(*this);
  }
  void setup(model* m) override;
  std::string name() const override { return "cosine decay learning rate"; }

protected:
  float global_schedule(model* m) override;
  float optimizer_schedule(model* m, optimizer& opt) override;

private:
  void write_specific_proto(lbann_data::Callback& proto) const final;

  float m_lr_max;
  float m_lr_min;
  size_t m_decay_steps;
  float m_initial_lr;
  size_t m_warmup_steps;
};

// Builder function
std::unique_ptr<callback_base>
build_cosine_decay_learning_rate_callback_from_pbuf(
  const google::protobuf::Message&,
  std::shared_ptr<lbann_summary> const&);

} // namespace callback
} // namespace lbann

#endif // LBANN_CALLBACKS_LEARNING_RATE_HPP_INCLUDED