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]

1. Fields can be either private or public, it really depends on the implementation and usage. The visibility of a field is up to what you expect users to do with it.
2. Rust is not an object-oriented language, so getters and setters are not as prominent. If you want to modify or retrieve the value of a private field, you could use something like a getter or setter. Otherwise, the visibility of the field means that you can access and/or modify it. It can either be public, public for the current crate only, or private.
3. Yep, there should be no issue with that 🙂 though these are treated as constant records so their values won't be serialized and thus have to be correctly initialized creating a new instance that you will load the record into.

[wangjiawen2013]

let model: MyModel<B> = MyModel {
    linear: LinearConfig::new(10, 16).init(&device),
    layer: MyLayer {
        l1: LinearConfig::new(16, 8).init(&device),
        l2: LinearConfig::new(8, 2).init(&device),
    },
};

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

I tested the above model.linear.weight.map(|x| x.clamp(-0.5, 0.5) and find the model was consumed and I cannot see the weight after clamping. It means that we cannot use the modified model later, which is a unexpected behavior.

[laggui]

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));

[wangjiawen2013]

Yes, it works !
Now if I want modify the weight with a tensor/array, how to operate ? I tried model.linear.weight.val() = my_tensor, it didn't work. There is no usage scenario for this operation. I am just curious about it and perhaps it will be used somewhere.

[laggui]

That's because param.val() returns a clone of the tensor held by the parameter, so you're not actually modifying the parameter tensor being held.

You could still use param.map(...) to achieve that, for example:

let value = Tensor::<B, 2>::ones(model.linear.weight.shape(), &device);

// Overwrite weight with value tensor
model.linear.weight = model.linear.weight.map(|_x| value);

[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!