fix a bug with ema thanks to Yoad, also allow for initializing concept embed id with superclass, as suggested in feedback

Author: lucidrains

perfusion_pytorch/embedding.py

@@ -25,7 +25,8 @@ class EmbeddingWrapper(Module):
     def __init__(
         self,
         embed: nn.Embedding,
-        num_concepts = 1
+        num_concepts = 1,
+        superclass_embed_id: Optional[Union[int, Tuple[int, ...]]] = None
     ):
         super().__init__()
         self.embed = embed
@@ -34,7 +35,20 @@ def __init__(
         self.num_embeds = num_embeds
         self.num_concepts = num_concepts
         self.concepts = nn.Parameter(torch.zeros(num_concepts, dim))
-        nn.init.normal_(self.concepts, std = 0.02)
+
+        if exists(superclass_embed_id):
+            # author had better results initializing the concept embed to the super class embed, allow for that option
+
+            if not isinstance(superclass_embed_id, tuple):
+                superclass_embed_id = (superclass_embed_id,)
+
+            superclass_embed_indices = torch.tensor(list(superclass_embed_id))
+            superclass_embeds = embed(superclass_embed_indices)
+            self.concepts.data.copy_(superclass_embeds)
+        else:
+            # otherwise initialize to usually small init for embeds
+
+            nn.init.normal_(self.concepts, std = 0.02)
 
         self.concept_embed_ids = tuple(range(num_embeds, num_embeds + num_concepts))
 
@@ -44,7 +58,7 @@ def parameters(self):
     def forward(
         self,
         x,
-        concept_id: Optional[Union[int, Tuple[int, ...]]] = None
+        concept_id: Optional[Union[int, Tuple[int, ...]]] = None,
     ):
         concept_masks = tuple(concept_id == x for concept_id in self.concept_embed_ids)
 

perfusion_pytorch/perfusion.py

@@ -316,7 +316,11 @@ def forward(
 
             if not initted:
                 self.initted[concept_id].data.copy_(Tensor([True]))
-                self.ema_concept_text_encs[concept_id].data.copy_(concept_text_enc)
+
+            # update ema i_*
+
+            self.ema_concept_text_encs[concept_id].data.copy_(concept_text_enc)
+
         else:
             assert self.initted[concept_id_tensor].all(), 'you have not initialized or trained this module for the concepts id given'
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'perfusion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.30',
+  version = '0.0.31',
   license='MIT',
   description = 'Perfusion - Pytorch',
   author = 'Phil Wang',