Upload 8 files

delta-iris/src/models/convnet.py

@@ -1,11 +1,7 @@
-from dataclasses import dataclass
-from typing import List
-
-from einops import rearrange
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
+from einops import rearrange
 
 class FrameEncoder(nn.Module):
     def __init__(self, config: dict) -> None:

delta-iris/src/models/quantizer.py

@@ -0,0 +1,48 @@
+import math
+from typing import Optional
+from einops import rearrange
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Quantizer(nn.Module):
+    def __init__(self, codebook_size: int, codebook_dim: int, input_dim: int, max_codebook_updates_with_revival: Optional[int] = None) -> None:
+        super().__init__()
+        assert math.log2(codebook_size).is_integer()
+        self.revival_entropy_threshold = int(math.log2(codebook_size)) - 2
+        self.max_codebook_updates_with_revival = max_codebook_updates_with_revival
+        self.pre_quant_proj = nn.Linear(input_dim, codebook_dim)
+        self.post_quant_proj = nn.Linear(codebook_dim, input_dim)
+        codebook = torch.empty(codebook_size, codebook_dim, requires_grad=False).uniform_(-1.0 / codebook_size, 1.0 / codebook_size)
+        self.num_codebook_updates = torch.tensor(0)
+        self.codebook = codebook
+        self.codewords_freqs = torch.ones(codebook_size).div(codebook_size)
+
+    def forward(self, z: torch.Tensor) -> dict:
+        z = self.pre_quant_proj(z)
+        z = F.normalize(z, dim=-1)
+        b, k = z.size(0), z.size(2)
+        z = rearrange(z, 'b t k e -> (b t k) e')
+
+        cosine_similarity = torch.einsum('n e, c e -> n c', z, self.codebook)
+        tokens = cosine_similarity.argmax(dim=-1)
+        q = self.codebook[tokens]
+
+        q = z + (q - z).detach()
+        q = self.post_quant_proj(q)
+
+        q = rearrange(q, '(b t k) e -> b t k e', b=b, k=k)
+        tokens = rearrange(tokens, '(b t k) -> b t k', b=b, k=k)
+        return {
+            "q": q,
+            "tokens": tokens,
+        }
+
+    def compute_codebook_entropy(self) -> float:
+        probs = self.codewords_freqs[self.codewords_freqs != 0]
+        return -(torch.log2(probs) * probs).sum().item()
+
+    @torch.no_grad()
+    def embed_tokens(self, tokens: torch.LongTensor) -> torch.FloatTensor:
+        return self.post_quant_proj(self.codebook[tokens])

delta-iris/src/models/transformer.py

@@ -2,10 +2,9 @@
 Inspired from https://github.com/karpathy/minGPT
 """
 
-from dataclasses import dataclass
 from typing import Optional
-
 from einops import rearrange
+
 import torch
 import torch.nn as nn
 

delta-iris/src/tokenizer.py

@@ -1,14 +1,10 @@
-from dataclasses import dataclass
 import math
-from typing import Dict, Tuple
-
 from einops import rearrange
 import torch
 import torch.nn as nn
 
 from .models.convnet import FrameEncoder, FrameDecoder
-from .models.tokenizer.quantizer import Quantizer
-from .models.utils import init_weights
+from .models.quantizer import Quantizer
 
 class Tokenizer(nn.Module):
     def __init__(self, config: dict) -> None:
@@ -32,8 +28,6 @@ class Tokenizer(nn.Module):
         self.decoder = FrameDecoder(config["decoder_config"])
         self.frame_cnn = FrameEncoder(config["frame_cnn_config"])
 
-        self.apply(init_weights)
-
     def __repr__(self) -> str:
         return "tokenizer"
 

delta-iris/src/world_model.py

@@ -1,15 +1,11 @@
-from dataclasses import dataclass
-
-from einops import rearrange, repeat
+from einops import rearrange
 from einops.layers.torch import Rearrange
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from .models.convnet import FrameEncoder
 from .models.slicer import  Head
 from .models.transformer import TransformerEncoder
-from .models.utils import init_weights
 
 class WorldModel(nn.Module):
     def __init__(self, config: dict) -> None:
@@ -55,8 +51,6 @@ class WorldModel(nn.Module):
             )
         )
 
-        self.apply(init_weights)
-
     def __repr__(self) -> str:
         return "world_model"