Custom initialization for nn layers?

[theunnecessarythings]

Do we have a way to do custom weight initialization for nn layers (conv, linear, norms etc). For eg. in tch-rs the configs had ws_init and bs_init. From what I have seen from the code, these are initialized locally in the functions/modules (for candle). Am I missing something?

[LaurentMazare]

We usually provide two ways to create nn layers, e.g. via conv1d that would perform the initialization if required or via Conv1d::new that would take the weight tensor as inputs. conv1d calls Conv1d::new and you can easily duplicate the related code to perform a different initialization if you want, e.g. just copy this function and tweak the "kaiming normal" bit to your liking.

[laptou]

Is there a way to do an initialization besides the ones in the Init struct? I want to initialize a parameter using linspace/arange, unless it already exists in the VarBuilder, but I can't figure out how to do that.

[LaurentMazare]

An easy way though not very efficient is to initialize it using one of the built-in initializers, and then set it to the value that you want (either in the same module or at the level of your training script).
A more complicated way would be to write your own var-builder to handle the new initialization but not sure it's worth the hassle.

[laptou]

But if I set it to a certain value, then won't that overwrite the value if it was present in the VarBuilder? Essentially, I want my vector to have a default value if I'm initializing my model to the default state, but this tensor's value can change over time. If I'm loading a pretrained version of this model then I want the trained value, not the original linspace default

[LaurentMazare]

Right, it would indeed overwrite the value assuming that it's a variable and not a tensor, but you can probably control from your code whether you want to do the set or not. It might be as simple as checking whether the tensor is a variable or not, a variable comes from a VarMap (used for training) and a tensor from the other varbuilders. You just have to be a bit cautious if you start loading some checkpoint (i.e. use a VarMap and get the initial values from a safetensor file).

[laptou]

Ah, I didn't know about Tensor::is_variable. Thanks!