Get weights and bias from the saved model

[wangjiawen2013]

Hi,
where we can find the weights and bias from the saved burn model ?

[laggui]

The parameters for each module are accessible via their respective field. For example, you can get the Linear weight and bias tensor values with:

let weight = linear.weight.val();
let bias = linear.bias.unwrap().val();

[wangjiawen2013]

@laggui
Thanks, and this is a simple case. However, things become complicated in my case.
I construct a model consisting of many blocks, and there are many various layers in each block, such as linear, conv2d and leakyrelu and so. And I want to clip the parameters.

In summary, in a complicated nested model, how to get and modify the parameters in a convenient way ?

With pytorch, this can be achieved by the following code:

 for p in discriminator.parameters():
    print(p.data.type)
    p.data.clamp(-clip_value, clip_value)

[laggui]

Ahhh ok in this case you probably want to use a ModuleMapper similar to what is implemented for the optimizer adapter which includes gradient clipping (though you're not dealing with gradients, just the values themselves in your case).

Each float parameter will go through this mapping so you can simply apply your clipping there.

[wangjiawen2013]

@laggui
Hi,
I am troubling you again!
I am trying to hard to under the working mechanism of ModuleMapper but failed to understand what it is. This is the definition of some methods:

/// Module visitor trait.
pub trait ModuleVisitor<B: Backend> {
    /// Visit a tensor in the module.
    fn visit<const D: usize>(&mut self, id: ParamId, tensor: &Tensor<B, D>);
}

/// Module mapper trait.
pub trait ModuleMapper<B: Backend> {
    /// Map a tensor in the module.
    fn map<const D: usize>(&mut self, id: ParamId, tensor: Tensor<B, D>) ->
      Tensor<B, D>;
}

1. Whait's the parameter tensor in in visit and map does ? And what's the tensor that map returns do ? is it different from the tensor in the input parameter ? I think we just want to visit a paramid, so why a tensor is needed in the visit parameter ?

2. How to get all the ParamId in a layerblock or model ?

3. Could you show a simple example on how to modify the parameters of a model ?

[laggui]

1. The tensor contains the floating-point tensor parameter, so if you want to modify its value (e.g., clip between a range) this is the tensor that you want to operate on. Since you want to modify its value, you should use the ModuleMapper as it allows you to take the original parameter tensor as input and return its modified value.
2. In order to get all the ParamId in a module you need to iterate over all the parameters to capture a tensor param's ParamId. If you don't care about mapping a specific parameter id to its module, then you could possibly implement the ModuleVisitor to fill a container with all the parameter ids of your module.
3. Here is a quick example which very conveniently clips all module parameters between a range of values.

use burn::{
    backend::NdArray,
    module::{Module, ModuleMapper},
    nn::{Linear, LinearConfig},
    prelude::Backend,
};

#[derive(Debug, Module)]
pub struct MyModel<B: Backend> {
    pub linear: Linear<B>,
}

pub struct Clip {
    pub min: f32,
    pub max: f32,
}

impl<B: Backend> ModuleMapper<B> for Clip {
    fn map_float<const D: usize>(
        &mut self,
        _id: burn::module::ParamId,
        tensor: burn::prelude::Tensor<B, D>,
    ) -> burn::prelude::Tensor<B, D> {
        tensor.clamp(self.min, self.max)
    }
}

fn main() {
    type B = NdArray;
    let device = Default::default();

    // Initial weights
    let model = MyModel {
        linear: LinearConfig::new(10, 2).init::<B>(&device),
    };
    println!("Before:\n{}\n", model.linear.weight.val());

    // Clip module mapper
    let mut clip = Clip {
        min: -0.01,
        max: 0.01,
    };
    let model = model.map(&mut clip);

    println!("After:\n{}\n", model.linear.weight.val());
}

[wangjiawen2013]

I reproduced this example successfully, however, when I use "clip" in my model, things became weired ! The parameters updated as expected when I didn't clip the parameters, but when I clip them, the parameters were not updated at all during the training !

This is the clip I used:

/// Clip module mapper to clip all module parameters between a range of values
#[derive(Module, Clone, Debug)]
pub struct Clip {
    pub min: f32,
    pub max: f32,
}

impl<B: Backend> ModuleMapper<B> for Clip {
    /// Map an float tensor in the module. Each ParamId and it's tensor value will be transfered
    /// to map_float and the modified value will be retured.
    fn map_float<const D: usize> (&mut self, _id: ParamId, tensor: Tensor<B, D>) -> Tensor <B, D> {
        tensor.clamp(self.min, self.max)
    }
}

And this is a chunk of code from the training.rs, the parameters were not updated at all !

let mut clip = Clip {min: -0.01, max: 0.01};

let loss_d = -discriminator.forward(batch.images).mean() + discriminator.forward(fake_images).mean();

// Gradients for the current backward pass
let grads = loss_d.backward();
// Gradients linked to each parameter of the discriminator
let grads = GradientsParams::from_grads(grads, &discriminator);
// Update the discriminator using the optimizer
discriminator = optimizer_d.step(config.lr, discriminator, grads);
// Clip parameters (weights) of discriminator
discriminator = discriminator.map(&mut clip);

When I commented out the last line, the parameters can be updated normally during the training iteration. I thinks this is perhaps the reason which caused the failure of my model (see this disscussion #2593)

[laggui]

Ahhh ok it was not clear what the weight clipping was supposed to be for, but in this case you need to make sure that the require_grad attribute is handled correctly. Otherwise the resulting parameter might not be tracked.

This is similar to the optimizer adapter that I linked earlier in the discussion:

impl<B: AutodiffBackend> ModuleMapper<B> for Clip {
    fn map_float<const D: usize>(&mut self, _id: ParamId, tensor: Tensor<B, D>) -> Tensor<B, D> {
        let is_require_grad = tensor.is_require_grad();

        let mut tensor = Tensor::from_inner(tensor.inner().clamp(self.min, self.max));

        if is_require_grad {
            tensor = tensor.require_grad();
        }
        tensor
    }
}

It should fix the issue you're currently observing with wgan where the loss does not seem to really fluctuate (probably because parameters were never updated after the first time they were clipped).

[wangjiawen2013]

I have re-implemented wgan in burn, the performance is comparable with that of pytorch. I'll pull a request later.

[wangjiawen2013]

@laggui
Thanks for your help ! I can clip parameters now.
You said "In order to get all the ParamId in a module you need to iterate over all the parameters to capture a tensor param's ParamId.". l read carefully on burn's docs and source code but still cannot find how to do it. So my question is:

1. How get all the paramid in a module.
2. How to modify a specific parameters (I don't need this but it may help later.).
   A chunk of exemple code is prefered !

[laggui]

Dropout and LayerNorm are separate modules and applied between Lstm layers usually. So that would be in the stacked implementation 🤔

[wangjiawen2013]

That's my bad, I should have been more clear. .map(...) always returns the mapped value. You need to assign the mapped value:

// Make sure the model is mutable
// let mut model = ...;

// Apply a function to the weights (here, clamp between -0.5 and 0.5)
model.linear.weight = model.linear.weight.map(|x| x.clamp(-0.5, 0.5));

Things get a bit tricky when I overwrite the bias because the bias is Some(Param), which is different from that of weight. This is how I did:

let bias = Tensor::<NdArray, 1>::ones([16], &device);
let bias = Param::initialized(model.linear.bias.clone().unwrap().id, bias);
model.linear.bias = model.linear.bias.map(|_x| bias);

Does it make sense, or do you have an easy way ?

[laggui]

Since the bias is wrapped in Option<...>, the first .map(...) acts on the inner value (the tensor param). So you can do two maps, for example:

model.linear.bias = model.linear.bias.map(|x| x.map(|x_i| Tensor::ones_like(&x_i)));

[wangjiawen2013]

I have a sophisticated implementation of Long Short-Term Memory (LSTM) networks in Burn, featuring state-of-the-art architectural enhancements and optimizations. This implementation includes stacked and bidirectional processing capabilities and advanced regularization techniques (layernorm and dropout) and combined weight matrices for efficiency, making it suitable for both research and production environments. More details can be found at the PyTorch implementation. Here is the test result:

This implementation is complementary to Burn's official LSTM, users can choose either one depends on the project's specific needs. Could I pull a request to burn's example?

[laggui]

This refers to the stacked LSTM previously discussed, right? Where multiple instances of the Lstm cell implemented in Burn are used to implement a multi-layer LSTM.

If so, I think this could help a lot of users that want to use this type of network for their task. We can start with a PR for the example, and possibly add it to the official nn modules in Burn if it can be useful.

[wangjiawen2013]

l read carefully on burn's docs and source code. And here is my suggestions:

1. There are too few tutorials on visitor & mapper in the burn's book (https://burn.dev/burn-book/building-blocks/module.html#visitor--mapper).
2. More advanced knowledge on rust language is needed (macros, closures), so it's hard for some users to understand the source code. For example, I cannot understand the definition of num_params and cannot even find the definition of map, I only find the signature of map (https://github.com/tracel-ai/burn/blob/main/crates/burn-core/src/module/base.rs). All these are necessary to operate on the parameters.
   That's why I meet so much difficulties on how to operate on the parametes. We're not as familiar with burn as you developers. So more tutorial or annotation is prefered in burn's book and source code, then we can find what we need easily (such as the way to get all the paramids).

[laggui]

1. That's a good point! Would adding the example I posted above to the book help?
2. Ahhh yes I noticed that the book did not provide the up to date trait, fixed it in PR Fix module visitor and mapper trait definition in the book #2609 😅

[wangjiawen2013]

Good!