Is there a way to dynamically modify methods in Rust?

myster 0
Viewed 94

I have a struct, Node, of which one of the fields (node_type) is an enum LayerType (which is either LayerType::ReLU, LayerType::Sigmoid, LayerType::Tanh or LayerType::Linear). There is a function on each Node: non_linear_function(&self, x: f64) -> f64 which should select the correct function to use based on the variant of the LayerType stored in node_type. The code I currently have is:

    pub fn non_linear_function(&self, x: f64) -> f64 {
        match self.node_type {
            LayerType::ReLU => todo!(),
            LayerType::Sigmoid => todo!(),
            LayerType::Tanh => todo!(),
            LayerType::Linear => todo!(), 
        }
    }

However, I am not sure whether or not this is efficient. Will it perform the check every time the function is called? If so, is there a way to prevent this happening?

All of the relevant code can be found below.

struct Layer {
    layer_type: LayerType,
    layer_width: usize,
    node_vec: Vec<Node>
}

#[derive(Debug)]
struct Node {
    prev_layer_size: usize,
    node_type: LayerType,
    weights: Vec<f64>,
    bias: f64
}

impl Node {
    pub fn new(prev_layer_size: usize, layer_pointer: &Layer) -> Node {
        // Instantiates a new node based on the size of the previous layer
        let mut rng = rand::thread_rng();
        return Node {
            prev_layer_size,
            node_type: (*layer_pointer).layer_type,
            weights: (0..prev_layer_size).map(|x| rng.gen()).collect(),
            bias: rng.gen()
        };
    }

    pub fn eval(&self, prev_activations: &Vec<f64>) -> f64 {
        // Evaluate the activation of the node 
        return self.non_linear_function(
            self.weights.iter().zip(prev_activations.iter()).map(
                |(&x, &y)| x*y
            ).sum::<f64>() + self.bias
        )
    }

    pub fn update(&mut self, target: f64) -> Vec<f64> {
        // Changes weights and biases based on target,
        // which indicates how much a given node should change and
        // returns a vector for how it thinks the nodes in the layer
        // behind it should change
        todo!()
    }

    // The code in question is here
    pub fn non_linear_function(&self, x: f64) -> f64 {
        match self.node_type {
            LayerType::ReLU => todo!(),
            LayerType::Sigmoid => todo!(),
            LayerType::Tanh => todo!(),
            LayerType::Linear => todo!(), 
        }
    }
}

#[derive(Debug, Clone, Copy)]
enum LayerType {
    ReLU,
    Sigmoid,
    Tanh,
    Linear
}

Is there an idiomatic way to implement this without match? Or is the match already optimised away by the compiler? Additionally, I can see that it is also an option to implement several different versions of Node, for each different LayerType, however this does not sound like a good idea. I also understand that this is probably a miniscule time save, but I am writing this code to learn how to write fast Rust code.

neural-network rust
2 Answers

Will it perform the check every time the function is called?

Yes, with this code the check will be kept at runtime.

If so, is there a way to prevent this happening?

Instead of using an enum for the LayerType, use a trait that specifies the method non_linear_function and implement it for the activation functions you want to support, this will move the check at compile time.


pub trait ActivationFunction {
    fn non_linear_function(&self, x: f64) -> f64;
}

struct ReLU;

impl ActivationFunction for ReLU {
    fn non_linear_function(&self, x: f64) -> f64 {
        f64::max(0.0, x)
    }
}

// ...

self.node_type.non_linear_function(x)

// ...
Ismail Maj0

We can see in your code that node_type in Node is probably redundant with layer_type in Layer: the non-linear function should then be chosen at the Layer level.

If your concern is efficiency, I suggest you step away from this OO-like design. Instead of asking each node of the layer « how would you do that? » and obtaining the same behaviour for each one, you could put all the logic in the layer and consider the nodes as simple passive structs (and even go further by packing properties as suggested in data-oriented-design).

The key is that the details of the computation are known once for all for the entire layer, then the compiler can generate some linear code with no branching/indirection.

In the following code, I tried to stay as close as possible to the example you provided. The eval() function in Node does not apply the non-linear function. Instead, I added an eval() function in Layer, which chooses once for all the non-linear function (I used dummy formulas) to apply to each node. The Node struct does not need to hold the LayerType since it is already held in the Layer. It does not either need to keep prev_layer_size since it is redundant with weights.len().

Note however that the inner-product of weight and prev_activations will likely take much more time than the match to choose the non-linear function to be applied right after; the efficiency gain expected on the way we specify this function may not be substantial in the end...

struct Layer {
    layer_type: LayerType,
    layer_width: usize,
    node_vec: Vec<Node>,
}

impl Layer {
    pub fn eval(
        &self,
        prev_activations: &Vec<f64>,
    ) -> Vec<f64> {
        let eval_nodes =
            self.node_vec.iter().map(|n| n.eval(prev_activations));
        match self.layer_type {
            LayerType::ReLU => {
                let fnct = |x| 2.0 * x; // dummy formula
                eval_nodes.map(fnct).collect()
            }
            LayerType::Sigmoid => {
                let fnct = |x| 2.0 * x; // dummy formula
                eval_nodes.map(fnct).collect()
            }
            LayerType::Tanh => {
                let fnct = |x| 2.0 * x; // dummy formula
                eval_nodes.map(fnct).collect()
            }
            LayerType::Linear => {
                let fnct = |x| 2.0 * x; // dummy formula
                eval_nodes.map(fnct).collect()
            }
        }
    }
}

#[derive(Debug)]
struct Node {
    weights: Vec<f64>,
    bias: f64,
}

impl Node {
    pub fn new(
        prev_layer_size: usize,
        layer_pointer: &Layer,
    ) -> Node {
        // Instantiates a new node based on the size of the previous layer
        let mut rng = rand::thread_rng();
        Node {
            weights: (0..prev_layer_size).map(|x| rng.gen()).collect(),
            bias: rng.gen(),
        }
    }

    pub fn eval(
        &self,
        prev_activations: &Vec<f64>,
    ) -> f64 {
        // Evaluate the activation of the node
        // NOTE: the non linear function is not used here
        self.weights
            .iter()
            .zip(prev_activations.iter())
            .map(|(&x, &y)| x * y)
            .sum::<f64>()
            + self.bias
    }

    pub fn update(
        &mut self,
        target: f64,
    ) -> Vec<f64> {
        // Changes weights and biases based on target,
        // which indicates how much a given node should change and
        // returns a vector for how it thinks the nodes in the layer
        // behind it should change
        todo!()
    }
}

#[derive(Debug, Clone, Copy)]
enum LayerType {
    ReLU,
    Sigmoid,
    Tanh,
    Linear,
}
prog-fh0
Related