TRF062

_init_weights must initialize tensors with the transformers initialization module, not torch.nn.init.

   
Default Enabled
Scope All models
Source mlinter/trf062.py
Show in terminal mlinter --rule TRF062

What it does

Flags calls inside _init_weights to an in-place torch.nn.init initializer such as nn.init.normal_(...), however it was imported (torch.nn.init, from torch.nn import init, from torch.nn.init import normal_). Helpers that only compute a value, such as calculate_gain or _calculate_fan_in_and_fan_out, are not flagged.

Why is this bad?

The transformers init functions skip a tensor whose _is_hf_initialized flag is set, so weights already loaded from a checkpoint are never overwritten. torch.nn.init has no such check. from_pretrained only covers for it by patching the torch.nn.init module while initialize_weights runs, which misses a function imported by name: from torch.nn.init import normal_ keeps the unpatched function, and the loaded weight is silently replaced by a random one. The transformers functions hold the guarantee themselves instead of relying on that patch.

Example

 from torch import nn
+
+from ... import initialization as init


 def _init_weights(self, module):
     super()._init_weights(module)
     if isinstance(module, AcmeLayerScale):
-        nn.init.constant_(module.lambda1, self.config.layer_scale_init_value)
+        init.constant_(module.lambda1, self.config.layer_scale_init_value)

Suppressing this rule

Add a # trf-ignore: TRF062 comment on the flagged line or the line directly above it. See Suppressing rules for whole-file directives and when a suppression is the wrong answer.