optional structure fields

[wangjiawen2013]

Hi,
Here is a chunk of pytorch code that can generate a block with or without batchnorm:

When I use burn instead, I find it's hard to implement similar function. How to set norm as an optional field ?

#[derive(Module, Debug)]
pub struct Block<B: Backend> {
    fc: nn::Linear<B>,
    norm: nn::BatchNorm<B, 1>,
    activation: nn::LeakyRelu,
}

impl<B: Backend> Block<B> {
    pub fn new(in_features: usize, out_features: usize, device: &B::Device) -> Self {
        let fc = nn::LinearConfig::new(in_features, out_features)
            .with_bias(True)
            .init(device);
        let norm = nn::BatchNormConfig::new(out_features, 0.8).init(device);
        let activation = nn::LeakyReluConfig::new().with_negative_slope(0.2).init(device);

        Self {
            fc,
            norm,
            activation
        }
    }
}

[laggui]

This can be achieved via a simple Option field.

Take this example from the ResNet implementation:

#[derive(Module, Debug)]
pub struct BasicBlock<B: Backend> {
    conv1: Conv2d<B>,
    bn1: BatchNorm<B, 2>,
    relu: Relu,
    conv2: Conv2d<B>,
    bn2: BatchNorm<B, 2>,
    downsample: Option<Downsample<B>>, // optional
}

impl<B: Backend> BasicBlock<B> {
    fn forward(&self, input: Tensor<B, 4>) -> Tensor<B, 4> {
        let identity = input.clone();

        let out = self.conv1.forward(input);
        let out = self.bn1.forward(out);
        let out = self.relu.forward(out);
        let out = self.conv2.forward(out);
        let out = self.bn2.forward(out);

        let out = {
           // Optional module
            match &self.downsample {
                Some(downsample) => out + downsample.forward(identity),
                None => out + identity,
            }
        };

        // Activation
        self.relu.forward(out)
    }
}