Properly initialize prototypes in LVQMLN

2021-05-09 20:55:28 +02:00 · 2021-05-09 20:55:28 +02:00 · c6e06ceaa4
commit c6e06ceaa4
parent 7a86bb19a8
1 changed files with 3 additions and 0 deletions
--- a/prototorch/models/glvq.py
+++ b/prototorch/models/glvq.py
@ -253,6 +253,9 @@ class LVQMLN(GLVQ):
                 **kwargs):
        super().__init__(hparams, **kwargs)
        self.backbone = backbone_module(**backbone_params)
+        with torch.no_grad():
+            protos = self.backbone(self.proto_layer()[0])
+        self.proto_layer.load_state_dict({"_components": protos}, strict=False)

    def forward(self, x):
        latent_protos, _ = self.proto_layer()