3min epoch breakout

https://wandb.ai/wassname/iris/runs/1zjsneyu?workspace=user-wassname

src/models/actor_critic.py

@@ -36,18 +36,19 @@ class ImagineOutput:
 class ActorCritic(nn.Module):
     def __init__(self, act_vocab_size, use_original_obs: bool = False) -> None:
         super().__init__()
-        shrink = 4
+        shrink = 8
+        s = 2
         self.use_original_obs = use_original_obs
-        self.conv1 = nn.Conv2d(3, 32, 3, stride=1, padding=1)
+        self.conv1 = nn.Conv2d(3, 32//s, 3, stride=1, padding=1)
         self.maxp1 = nn.MaxPool2d(2, 2)
-        self.conv2 = nn.Conv2d(32, 32, 3, stride=1, padding=1)
+        self.conv2 = nn.Conv2d(32//s, 32//s, 3, stride=1, padding=1)
         self.maxp2 = nn.MaxPool2d(2, 2)
-        self.conv3 = nn.Conv2d(32, 64, 3, stride=1, padding=1)
+        self.conv3 = nn.Conv2d(32//s, 64//s, 3, stride=1, padding=1)
         self.maxp3 = nn.MaxPool2d(2, 2)
-        self.conv4 = nn.Conv2d(64, 64//shrink, 3, stride=1, padding=1)
+        self.conv4 = nn.Conv2d(64//s, 64//shrink, 3, stride=1, padding=1)
         self.maxp4 = nn.MaxPool2d(2, 2)
 
-        self.lstm_dim = 64
+        self.lstm_dim = 16
         self.lstm = nn.LSTMCell(1024//shrink, self.lstm_dim)
         self.hx, self.cx = None, None
 

src/models/transformer.py

@@ -17,94 +17,23 @@ from .kv_caching import KeysValues, KVCache
 @dataclass
 class TransformerConfig:
 
-        
-    tokens_per_block: int
     max_blocks: int
-    attention: str
 
     num_layers: int
     num_heads: int
     embed_dim: int
-
-    embed_pdrop: float
-    resid_pdrop: float
-    attn_pdrop: float
+    tokens_per_block: int
     
     # model_name: str = "stabilityai/stablelm-3b-4e1t"
     # https://huggingface.co/PY007/TinyLlama-1.1B-intermediate-step-715k-1.5T
     model_name: str = "PY007/TinyLlama-1.1B-intermediate-step-715k-1.5T"
     dropout: float = 0
     rank: int = 32
-    z_dim: int = 768
-    start_from: str = None
 
     @property
     def max_tokens(self):
         return self.tokens_per_block * self.max_blocks
 
-class Block(nn.Module):
-    def __init__(self, config: TransformerConfig) -> None:
-        super().__init__()
-        self.ln1 = nn.LayerNorm(config.embed_dim)
-        self.ln2 = nn.LayerNorm(config.embed_dim)
-        self.attn = SelfAttention(config)
-        self.mlp = nn.Sequential(
-            nn.Linear(config.embed_dim, 4 * config.embed_dim),
-            nn.GELU(),
-            nn.Linear(4 * config.embed_dim, config.embed_dim),
-            nn.Dropout(config.resid_pdrop),
-        )
-
-    def forward(self, x: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
-        x_attn = self.attn(self.ln1(x), past_keys_values)
-        x = x + x_attn
-        x = x + self.mlp(self.ln2(x))
-        return x
-
-
-class SelfAttention(nn.Module):
-    def __init__(self, config: TransformerConfig) -> None:
-        super().__init__()
-        assert config.embed_dim % config.num_heads == 0
-        assert config.attention in ('causal', 'block_causal')
-        self.num_heads = config.num_heads
-        self.key = nn.Linear(config.embed_dim, config.embed_dim)
-        self.query = nn.Linear(config.embed_dim, config.embed_dim)
-        self.value = nn.Linear(config.embed_dim, config.embed_dim)
-        self.attn_drop = nn.Dropout(config.attn_pdrop)
-        self.resid_drop = nn.Dropout(config.resid_pdrop)
-        self.proj = nn.Linear(config.embed_dim, config.embed_dim)
-
-        causal_mask = torch.tril(torch.ones(config.max_tokens, config.max_tokens))
-        block_causal_mask = torch.max(causal_mask, torch.block_diag(*[torch.ones(config.tokens_per_block, config.tokens_per_block) for _ in range(config.max_blocks)]))
-        self.register_buffer('mask', causal_mask if config.attention == 'causal' else block_causal_mask)
-
-    def forward(self, x: torch.Tensor, kv_cache: Optional[KVCache] = None) -> torch.Tensor:
-        B, T, C = x.size()
-        if kv_cache is not None:
-            b, nh, L, c = kv_cache.shape
-            assert nh == self.num_heads and b == B and c * nh == C
-        else:
-            L = 0
-
-        q = self.query(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)   # (B, nh, T, hs)
-        k = self.key(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)     # (B, nh, T, hs)
-        v = self.value(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)   # (B, nh, T, hs)
-
-        if kv_cache is not None:
-            kv_cache.update(k, v)
-            k, v = kv_cache.get()
-
-        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
-        att = att.masked_fill(self.mask[L:L + T, :L + T] == 0, float('-inf'))
-        att = F.softmax(att, dim=-1)
-        att = self.attn_drop(att)
-        y = att @ v
-        y = rearrange(y, 'b h t e -> b t (h e)')
-
-        y = self.resid_drop(self.proj(y))
-
-        return y
 
 class Transformer(nn.Module):
     def __init__(self, config: TransformerConfig) -> None: