rnn.h

#pragma once

#include <torch/nn/cloneable.h>
#include <torch/nn/modules/dropout.h>
#include <torch/nn/pimpl.h>
#include <torch/types.h>

#include <ATen/ATen.h>
#include <c10/util/Exception.h>

#include <cstddef>
#include <functional>
#include <memory>
#include <vector>

namespace torch {
namespace nn {

struct TORCH_API RNNOutput {
  Tensor output;
  Tensor state;
};

namespace detail {

struct TORCH_API RNNOptionsBase {
  RNNOptionsBase(int64_t input_size, int64_t hidden_size);
  virtual ~RNNOptionsBase() = default;
  TORCH_ARG(int64_t, input_size);
  TORCH_ARG(int64_t, hidden_size);
  TORCH_ARG(int64_t, layers) = 1;
  TORCH_ARG(bool, with_bias) = true;
  TORCH_ARG(double, dropout) = 0.0;
  TORCH_ARG(bool, bidirectional) = false;
  TORCH_ARG(bool, batch_first) = false;
};

template <typename Derived>
class RNNImplBase : public torch::nn::Cloneable<Derived> {
 public:
  enum class CuDNNMode { RNN_RELU = 0, RNN_TANH = 1, LSTM = 2, GRU = 3 };

  explicit RNNImplBase(
      const RNNOptionsBase& options_,
      optional<CuDNNMode> cudnn_mode = nullopt,
      int64_t number_of_gates = 1);

  void reset() override;

  void to(torch::Device device, torch::Dtype dtype, bool non_blocking = false)
      override;
  void to(torch::Dtype dtype, bool non_blocking = false) override;
  void to(torch::Device device, bool non_blocking = false) override;

  void pretty_print(std::ostream& stream) const override;

  void flatten_parameters();

  RNNOptionsBase options;

  std::vector<Tensor> w_ih;
  std::vector<Tensor> w_hh;
  std::vector<Tensor> b_ih;
  std::vector<Tensor> b_hh;

 protected:
  using RNNFunctionSignature = std::tuple<Tensor, Tensor>(
      /*input=*/const Tensor&,
      /*state=*/const Tensor&,
      /*params=*/TensorList,
      /*has_biases=*/bool,
      /*layers=*/int64_t,
      /*dropout=*/double,
      /*train=*/bool,
      /*bidirectional=*/bool,
      /*batch_first=*/bool);

  RNNOutput generic_forward(
      std::function<RNNFunctionSignature> function,
      const Tensor& input,
      Tensor state);

  std::vector<Tensor> flat_weights() const;

  bool any_parameters_alias() const;

  int64_t number_of_gates_;

  optional<CuDNNMode> cudnn_mode_;

  std::vector<Tensor> flat_weights_;
};
} // namespace detail

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ RNN ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

enum class RNNActivation : uint32_t {ReLU, Tanh};

struct TORCH_API RNNOptions {
  RNNOptions(int64_t input_size, int64_t hidden_size);

  RNNOptions& tanh();
  RNNOptions& relu();

  TORCH_ARG(int64_t, input_size);
  TORCH_ARG(int64_t, hidden_size);
  TORCH_ARG(int64_t, layers) = 1;
  TORCH_ARG(bool, with_bias) = true;
  TORCH_ARG(double, dropout) = 0.0;
  TORCH_ARG(bool, bidirectional) = false;
  TORCH_ARG(bool, batch_first) = false;
  TORCH_ARG(RNNActivation, activation) = RNNActivation::ReLU;
};

class TORCH_API RNNImpl : public detail::RNNImplBase<RNNImpl> {
 public:
  RNNImpl(int64_t input_size, int64_t hidden_size)
      : RNNImpl(RNNOptions(input_size, hidden_size)) {}
  explicit RNNImpl(const RNNOptions& options);

  void pretty_print(std::ostream& stream) const override;

  RNNOutput forward(const Tensor& input, Tensor state = {});

  RNNOptions options;
};

TORCH_MODULE(RNN);

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ LSTM ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

using LSTMOptions = detail::RNNOptionsBase;

class TORCH_API LSTMImpl : public detail::RNNImplBase<LSTMImpl> {
 public:
  LSTMImpl(int64_t input_size, int64_t hidden_size)
      : LSTMImpl(LSTMOptions(input_size, hidden_size)) {}
  explicit LSTMImpl(const LSTMOptions& options);

  RNNOutput forward(const Tensor& input, Tensor state = {});
};

TORCH_MODULE(LSTM);

// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ GRU ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

using GRUOptions = detail::RNNOptionsBase;

class TORCH_API GRUImpl : public detail::RNNImplBase<GRUImpl> {
 public:
  GRUImpl(int64_t input_size, int64_t hidden_size)
      : GRUImpl(GRUOptions(input_size, hidden_size)) {}
  explicit GRUImpl(const GRUOptions& options);

  RNNOutput forward(const Tensor& input, Tensor state = {});
};

TORCH_MODULE(GRU);

} // namespace nn
} // namespace torch