1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105

//! Definition of interfaces of neural networks.
use anyhow::Result;
use std::path::Path;
use tch::{nn, nn::VarStore, Tensor};

/// Base interface.
pub trait ModelBase {
    /// Trains the network given a loss.
    fn backward_step(&mut self, loss: &Tensor);

    /// Returns `var_store` as mutable reference.
    fn get_var_store_mut(&mut self) -> &mut nn::VarStore;

    /// Returns `var_store`.
    fn get_var_store(&self) -> &nn::VarStore;

    /// Save parameters of the neural network.
    fn save<T: AsRef<Path>>(&self, path: T) -> Result<()>;

    /// Load parameters of the neural network.
    fn load<T: AsRef<Path>>(&mut self, path: T) -> Result<()>;
}

/// Neural networks with a single input and a single output.
pub trait Model1: ModelBase {
    /// The input of the neural network.
    type Input;
    /// The output of the neural network.
    type Output;

    /// Performs forward computation given an input.
    fn forward(&self, xs: &Self::Input) -> Self::Output;

    /// TODO: check places this method is used in code.
    fn in_shape(&self) -> &[usize];

    /// TODO: check places this method is used in code.
    fn out_dim(&self) -> usize;
}

/// Neural networks with double inputs and a single output.
pub trait Model2: ModelBase {
    /// An input of the neural network.
    type Input1;
    /// The other input of the neural network.
    type Input2;
    /// The output of the neural network.
    type Output;

    /// Performs forward computation given a pair of inputs.
    fn forward(&self, x1s: &Self::Input1, x2s: &Self::Input2) -> Self::Output;
}

/// Neural network model that can be initialized with [VarStore] and configuration.
///
/// The purpose of this trait is for modularity of neural network models.
/// Modules, which consists a neural network, should share [VarStore].
/// To do this, structs implementing this trait can be initialized with a given [VarStore].
/// This trait also provide the ability to clone with a given [VarStore].
/// The ability is useful when creating a target network, used in recent deep learning algorithms in common.
pub trait SubModel {
    /// Configuration from which [SubModel] is constructed.
    type Config;

    /// Input of the [SubModel].
    type Input;

    /// Output of the [SubModel].
    type Output;

    /// Builds [SubModel] with [VarStore] and [SubModel::Config].
    fn build(var_store: &VarStore, config: Self::Config) -> Self;

    /// Clones [SubModel] with [VarStore].
    fn clone_with_var_store(&self, var_store: &VarStore) -> Self;

    /// A generalized forward function.
    fn forward(&self, input: &Self::Input) -> Self::Output;
}

/// Neural network model that can be initialized with [VarStore] and configuration.
///
/// The difference from [SubModel] is that this trait takes two inputs.
pub trait SubModel2 {
    /// Configuration from which [SubModel2] is constructed.
    type Config;

    /// Input of the [SubModel2].
    type Input1;

    /// Input of the [SubModel2].
    type Input2;

    /// Output of the [SubModel2].
    type Output;

    /// Builds [SubModel2] with [VarStore] and [SubModel2::Config].
    fn build(var_store: &VarStore, config: Self::Config) -> Self;

    /// Clones [SubModel2] with [VarStore].
    fn clone_with_var_store(&self, var_store: &VarStore) -> Self;

    /// A generalized forward function.
    fn forward(&self, input1: &Self::Input1, input2: &Self::Input2) -> Self::Output;
}