superclass does not need ema

Author: lucidrains

perfusion_pytorch/perfusion.py

@@ -88,7 +88,7 @@ def __init__(
 
         self.register_buffer('initted', torch.zeros(num_finetune_prompts).bool())
         self.register_buffer('ema_concept_text_encs', torch.zeros(num_finetune_prompts, dim_input))
-        self.register_buffer('ema_superclass_text_encs', torch.zeros(num_finetune_prompts, dim_input))
+        self.register_buffer('superclass_text_encs', torch.zeros(num_finetune_prompts, dim_input))
         self.register_buffer('superclass_outputs', torch.zeros(num_finetune_prompts, dim_output))
 
         # C in the paper, inverse precomputed
@@ -154,7 +154,12 @@ def forward(
             all_initted = initted.all()
 
             ema_concept_text_enc = self.ema_concept_text_encs[prompt_ids]
-            ema_superclass_text_enc = self.ema_superclass_text_encs[prompt_ids]
+
+            # fetch superclass
+
+            assert exists(superclass_text_enc) or all_initted
+
+            stored_superclass_text_enc = self.superclass_text_encs[prompt_ids]
 
             # for keys, the superclass output (o*) is stored on init
             # and never optimized
@@ -170,9 +175,9 @@ def forward(
                     concept_text_enc
                 )
 
-                ema_superclass_text_enc = torch.where(
+                superclass_text_enc = torch.where(
                     initted,
-                    ema_superclass_text_enc,
+                    stored_superclass_text_enc,
                     superclass_text_enc
                 )
 
@@ -182,17 +187,16 @@ def forward(
                     superclass_output
                 )
 
-            # exponential moving average of both concept and superclass
+            # exponential moving average for concept input encoding
 
             concept_text_enc = ema_concept_text_enc * decay + concept_text_enc * (1. - decay)
-            superclass_text_enc = ema_superclass_text_enc * decay + superclass_text_enc * (1. - decay)
 
             # store
 
             if not all_initted:
                 self.initted[prompt_ids] = True
                 self.ema_concept_text_encs[prompt_ids] = ema_concept_text_enc
-                self.ema_superclass_text_encs[prompt_ids] = ema_superclass_text_enc
+                self.superclass_text_encs[prompt_ids] = superclass_text_enc
                 self.superclass_outputs[prompt_ids] = superclass_output
 
         # take care of the output

setup.py

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