ASAC-pytorch 0.0.4__tar.gz → 0.0.7__tar.gz

{asac_pytorch-0.0.4 → asac_pytorch-0.0.7}/ASAC/ASAC.py

@@ -117,7 +117,7 @@ class Attention(Module):
             sim, indices, aux_loss, aux_loss_breakdown = self.attn_schema(orig_sim)
 
         if self.attn_add_residual:
-            sim = sim + orig_sim
+            sim = (sim + orig_sim) * 0.5
 
         # modulate
 
@@ -164,14 +164,14 @@ class AttentionSchema(Module):
         super().__init__()
 
         if not exists(encoder):
-            encoder = MLP(dim, dim_bottleneck, activation = nn.LeakyReLU())
+            encoder = MLP(dim, dim_bottleneck, dim_bottleneck, activation = nn.LeakyReLU())
 
         self.encoder = encoder
 
         self.vq = VectorQuantize(dim_bottleneck, **vq_kwargs)
 
         if not exists(decoder):
-            decoder = MLP(dim_bottleneck, dim, activation = nn.LeakyReLU())
+            decoder = MLP(dim_bottleneck, dim_bottleneck, dim, activation = nn.LeakyReLU())
 
         self.decoder = decoder
 
@@ -181,6 +181,8 @@ class AttentionSchema(Module):
         self.recon_loss_weight = recon_loss_weight
         self.commit_loss_weight = commit_loss_weight
 
+        self.register_buffer('zero', tensor(0.), persistent = False)
+
     def forward(
         self,
         attn_sim,
@@ -200,6 +202,8 @@ class AttentionSchema(Module):
 
         # loss, mse as in paper or reverse kl
 
+        recon_loss = self.zero
+
         if return_loss:
             if self.detach_target:
                 attn_sim = attn_sim.detach()
@@ -208,8 +212,8 @@ class AttentionSchema(Module):
                 recon_loss = F.kl_div(
                     attn_sim.log_softmax(dim = -1),
                     recon.softmax(dim = -1),
-                    reduction = 'batchmean'
-                )
+                    reduction = 'none'
+                ).sum(dim = -1).mean()
             else:
                 recon_loss = F.mse_loss(recon, attn_sim)
 
@@ -236,7 +240,8 @@ class ASAC(Module):
         dim_bottleneck = 256,
         vq_codebook_size = 256,
         recon_loss_weight = 1.,
-        commit_loss_weight = 1.
+        commit_loss_weight = 1.,
+        kl_div_loss = True
     ):
         super().__init__()
 
@@ -253,7 +258,8 @@ class ASAC(Module):
                 dim_bottleneck = dim_bottleneck,
                 codebook_size = vq_codebook_size,
                 recon_loss_weight = recon_loss_weight,
-                commit_loss_weight = commit_loss_weight
+                commit_loss_weight = commit_loss_weight,
+                kl_div_loss = kl_div_loss
             ) if use_asac and exists(seq_len) else None
 
             self.layers.append(ModuleList([
@@ -272,9 +278,7 @@ class ASAC(Module):
         if exists(self.pos_embedding):
             x = x + self.pos_embedding
 
-        total_aux_loss = 0.
-        total_recon_loss = 0.
-        total_commit_loss = 0.
+        total_aux_loss = total_recon_loss = total_commit_loss = 0.
 
         for attn, ff in self.layers:
             attn_out, indices, aux_loss, (recon_loss, commit_loss) = attn(x)

{asac_pytorch-0.0.4 → asac_pytorch-0.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ASAC-pytorch
-Version: 0.0.4
+Version: 0.0.7
 Summary: Implementation of Attention Schema-based Attention Control (ASAC)
 Project-URL: Homepage, https://pypi.org/project/ASAC/
 Project-URL: Repository, https://codeberg.org/lucidrains/ASAC
@@ -38,7 +38,7 @@ Requires-Dist: einops>=0.8.1
 Requires-Dist: einx>=0.3.0
 Requires-Dist: ema-pytorch
 Requires-Dist: torch-einops-utils>=0.1.2
-Requires-Dist: torch>=2.5
+Requires-Dist: torch>=2.4
 Requires-Dist: vector-quantize-pytorch
 Requires-Dist: x-mlps-pytorch
 Requires-Dist: x-transformers

{asac_pytorch-0.0.4 → asac_pytorch-0.0.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ASAC-pytorch"
-version = "0.0.4"
+version = "0.0.7"
 description = "Implementation of Attention Schema-based Attention Control (ASAC)"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -27,7 +27,7 @@ dependencies = [
     "einx>=0.3.0",
     "einops>=0.8.1",
     "ema-pytorch",
-    "torch>=2.5",
+    "torch>=2.4",
     "torch-einops-utils>=0.1.2",
     "vector-quantize-pytorch",
     "x-transformers",