How can this CNN be detifined in Burn?

[AutumnRoom]

Pytorch:
self.cnn = nn.ModuleList()
for _ in range(depth):
self.cnn.append(nn.Sequential(
weight_norm(nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding)),
LayerNorm(channels),
actv,
nn.Dropout(0.2),
))

here are some confusion,
1.weight_norm()
2.Sequential() and ModuleList
3.for _ in range(): //Maybe

[laggui]

1. We don't have any official API for weight reparameterization (see issue Support weight reparametrization #2738)
2. You can use a Vec of modules, though you will need to implement the forward pass for the loop. A vec contains elements of the same type so if you want to use different modules, you can use an enum. In your case, it is much easier (and cleaner in my opinion) to define a module block since it is always the same:

#[derive(Module, Debug, Clone)]
pub enum Activation {
    Relu,
    // etc.
}

#[derive(Module, Debug)]
pub struct ConvBlock<B: Backend> {
    pub conv: Conv1d<B>,
    pub norm: LayerNorm<B>,
    pub activation: Activation,
    pub dropout: Dropout,
}

impl<B: Backend> ConvBlock<B> {
    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
        let x = self.norm.forward(self.conv.forward(x));
        
        let x = match self.activation {
            Activation::Relu => burn::tensor::activation::relu(x),
        };

        self.dropout.forward(x)
    }
}

so your cnn would simply be a Vec<ConvBlock<B>>.
3. Given the previous suggestion, the number of layers in your vec can be defined when creating your module. You could define a config struct with an init(depth) method that returns the module initialized on the backend.

[laggui]

I've noticed that the "remove" and "register" methods are defined in burn::tensor::container::TensorContainer. I guess these two methods can be used to address the flaws in the previously proposed solution

Hmmm I'm not sure I understand how that will be useful. Maybe I missed something.

But I'm still a bit confused. How should I call the "register" and "remove" methods?

The TensorContainer is only used internally to keep the gradients during the backward pass. See for example how it is used in the optimizer adaptor (which is implementes the module mapper trait, actually).

During the back - propagation process of subsequent training, PyTorch's automatic differentiation mechanism will automatically calculate the gradients of weight_v and weight_g (weight components) to update the values of these components

We also have autodiff! So if you implement a solution to use the reparametrized weights (and discard the actual original weights), this will be automatically computed as well. But we currently don't have a straightforward solution to basically override a module's weights similar to what pytorch reparameterization does (if I understood that correctly).

[AutumnRoom]

I think you should understand it correctly.Burn doesn't have the parameter deconstruction operation like PyTorch, but I've written the following code.I wrote a WnormConv1d by referring to Conv1d in burn_core::nn::conv. I'm still a bit confused about parameter initialization. How can I use the KaimingUniform initialization method to initialize the weight component weight_v, and then initialize all the parameters of the weight component weight_g to 1? My biggest confusion is whether this WnormConv1d is the same as the Conv1d processed by the weight_norm function in PyTorch. There are many differences in details. However, I'd rather believe that the new module I constructed has skipped a lot of unnecessary steps and perfectly achieved the same effect.

#[derive(Config, Debug)]
pub struct WnormConv1dConfig {
    // etc.
}

#[derive(Module, Debug)]
#[module(custom_display)]
pub struct WnormConv1d<B: Backend> {
    pub weight_g: Param<Tensor<B, 3>>,
    pub weight_v: Param<Tensor<B, 3>>,
    // etc.
}

impl<B: Backend> ModuleDisplay for WnormConv1d<B> {
    // etc.
}

impl WnormConv1dConfig {
    pub fn init<B: Backend>(&self, device: &B::Device) -> WnormConv1d<B> {
        // etc.

        let shape_g = [
            self.channels_out,1,1
        ];
        let shape_v = [
            self.channels_out,
            self.channels_in / self.groups,
            self.kernel_size,
        ];

        // etc.
        let weight_g = self
            .initializer
            .init_with(shape_g, Some(fan_in), None, device);
        let weight_v = self
            .initializer
            .init_with(shape_v, Some(fan_in), None, device);

        // etc.

        WnormConv1d {
            weight_g,
            weight_v,
            // etc.
        }
    }
}

impl<B: Backend> WnormConv1d<B> {
    pub fn forward(&self, input: Tensor<B, 3>,dim: usize) -> Tensor<B, 3> {
        // etc.

        conv1d(
            input,
            self.weight_g.val() * (self.weight_v.val() / (self.weight_v.val()*self.weight_v.val()).sum_dim(dim).sqrt()),
            self.bias.as_ref().map(|bias| bias.val()),
            ConvOptions::new([self.stride], [padding], [self.dilation], self.groups),
        )
    }
}

[laggui]

However, I'd rather believe that the new module I constructed has skipped a lot of unnecessary steps and perfectly achieved the same effect.

If you only need to apply this to a single module type, then it is always going to be more straightforward. I think the forward definition is exactly how it should be! Keep the g and v parameters only, and compute the conv weight from those on-the-fly. Technically, this could be done for the bias as well, right? Or is it only applied to the conv weights?

With pytorch, there is a lot more additional complexity to overwrite the initial module parameters and I think they use forward/backward hooks. Plus it should support any module type. Bunch of abstractions that complicate the nature of the implementation.

I'm still a bit confused about parameter initialization. How can I use the KaimingUniform initialization method to initialize the weight component weight_v, and then initialize all the parameters of the weight component weight_g to 1?

In this case, you need to use different initializers for both. The current implementations in Burn use a single initializer for a module's weight, but you can do this with more granularity. You can split initializer into initializer_v and initializer_g in your WnormConv1dConfig or simply "hardcode" their usage when initializing the module. So you would have Initializer::KaimingUniform for one and Initializer::Ones for the other.

[AutumnRoom]

I set the dependency path of Burn to the local directory, wrote the modified WnormConv1d into the burn_core::nn::conv folder, and then added WnormConv1d to the project file. There was no error, and it seemed to be fine. However, some default settings may be required, and further adjustments may be needed in the future. weight_norm only acts on the weight of the module, and there is no need to quantize the bias. Only the weight components that the optimizer can act on are needed. In the future, it may be possible to decompose the weights of more modules and write new WnormConv2d, WnormConv3d, and WnormLinear.

[laggui]

Instead of using a local version of burn you can simply define this in your project! Not sure it would make it into burn_core::nn's official modules with this implementation since it will require a lot of duplication for different modules, but that's something that can be tackled in the future if we want to add this feature.

In the meantime, I think this should work great for your use case!