[FEAT][ImgPatchEmbed] [chore][disable_warnings_and_logs]

tests/nn/modules/test_img_patch_embed.py

@@ -0,0 +1,76 @@
+# FILEPATH: /Users/defalt/Desktop/Athena/research/zeta/tests/nn/modules/test_img_patch_embed.py
+
+import pytest
+from torch import nn
+import torch
+from zeta.nn.modules.img_patch_embed import ImgPatchEmbed
+
+
+def test_class_init():
+ model = ImgPatchEmbed()
+
+ assert isinstance(model.proj, nn.Conv2d)
+ assert model.img_size == 224
+ assert model.patch_size == 16
+ assert model.num_patches == 196
+
+
+def test_class_init_with_args():
+ model = ImgPatchEmbed(
+ img_size=448, patch_size=32, in_chans=1, embed_dim=512
+ )
+
+ assert isinstance(model.proj, nn.Conv2d)
+ assert model.img_size == 448
+ assert model.patch_size == 32
+ assert model.num_patches == 196
+ assert model.proj.in_channels == 1
+ assert model.proj.out_channels == 512
+
+
+def test_forward():
+ model = ImgPatchEmbed()
+ x = torch.randn(1, 3, 224, 224)
+ out = model(x)
+
+ assert out.shape == torch.Size([1, 196, 768])
+
+
+def test_forward_with_different_input():
+ model = ImgPatchEmbed()
+ x = torch.randn(2, 3, 224, 224)
+ out = model(x)
+
+ assert out.shape == torch.Size([2, 196, 768])
+
+
+def test_forward_with_different_img_size():
+ model = ImgPatchEmbed(img_size=448)
+ x = torch.randn(1, 3, 448, 448)
+ out = model(x)
+
+ assert out.shape == torch.Size([1, 196, 768])
+
+
+def test_forward_with_different_patch_size():
+ model = ImgPatchEmbed(patch_size=32)
+ x = torch.randn(1, 3, 224, 224)
+ out = model(x)
+
+ assert out.shape == torch.Size([1, 49, 768])
+
+
+def test_forward_with_different_in_chans():
+ model = ImgPatchEmbed(in_chans=1)
+ x = torch.randn(1, 1, 224, 224)
+ out = model(x)
+
+ assert out.shape == torch.Size([1, 196, 768])
+
+
+def test_forward_with_different_embed_dim():
+ model = ImgPatchEmbed(embed_dim=512)
+ x = torch.randn(1, 3, 224, 224)
+ out = model(x)
+
+ assert out.shape == torch.Size([1, 196, 512])

tests/nn/modules/test_simple_mamba.py

@@ -5,6 +5,7 @@
 from torch import nn
 from zeta.nn.modules.simple_mamba import Mamba, ResidualBlock, RMSNorm
 
+
 def test_mamba_class_init():
  model = Mamba(10000, 512, 6)
 
@@ -13,13 +14,15 @@ def test_mamba_class_init():
  assert isinstance(model.norm_f, RMSNorm)
  assert isinstance(model.lm_head, nn.Linear)
 
+
 def test_mamba_forward():
  model = Mamba(10000, 512, 6)
  x = torch.randint(0, 10000, (1, 50))
  out = model(x)
 
  assert out.shape == torch.Size([1, 50, 10000])
 
+
 def test_residual_block_class_init():
  block = ResidualBlock(512)
 
@@ -28,55 +31,63 @@ def test_residual_block_class_init():
  assert isinstance(block.fc1, nn.Linear)
  assert isinstance(block.fc2, nn.Linear)
 
+
 def test_residual_block_forward():
  block = ResidualBlock(512)
  x = torch.randn(1, 50, 512)
  out = block(x)
 
  assert out.shape == torch.Size([1, 50, 512])
 
+
 def test_mamba_different_vocab_size():
  model = Mamba(20000, 512, 6)
  x = torch.randint(0, 20000, (1, 50))
  out = model(x)
 
  assert out.shape == torch.Size([1, 50, 20000])
 
+
 def test_mamba_different_dim():
  model = Mamba(10000, 1024, 6)
  x = torch.randint(0, 10000, (1, 50))
  out = model(x)
 
  assert out.shape == torch.Size([1, 50, 10000])
 
+
 def test_mamba_different_depth():
  model = Mamba(10000, 512, 12)
  x = torch.randint(0, 10000, (1, 50))
  out = model(x)
 
  assert out.shape == torch.Size([1, 50, 10000])
 
+
 def test_residual_block_different_dim():
  block = ResidualBlock(1024)
  x = torch.randn(1, 50, 1024)
  out = block(x)
 
  assert out.shape == torch.Size([1, 50, 1024])
 
+
 def test_mamba_with_dropout():
  model = Mamba(10000, 512, 6, dropout=0.5)
  x = torch.randint(0, 10000, (1, 50))
  out = model(x)
 
  assert out.shape == torch.Size([1, 50, 10000])
 
+
 def test_residual_block_with_dropout():
  block = ResidualBlock(512, dropout=0.5)
  x = torch.randn(1, 50, 512)
  out = block(x)
 
  assert out.shape == torch.Size([1, 50, 512])
 
+
 def test_mamba_with_custom_layer():
  class CustomLayer(nn.Module):
  def forward(self, x):
@@ -86,4 +97,4 @@ def forward(self, x):
  x = torch.randint(0, 10000, (1, 50))
  out = model(x)
 
- assert out.shape == torch.Size([1, 50, 10000])
+ assert out.shape == torch.Size([1, 50, 10000])

zeta/nn/biases/relative_position_bias.py

@@ -6,6 +6,7 @@
 import torch
 from torch import nn
 
+
 class RelativePositionBias(nn.Module):
  def __init__(
  self,

zeta/nn/modules/__init__.py

@@ -46,6 +46,7 @@
 from zeta.nn.modules.visual_expert import VisualExpert
 from zeta.nn.modules.yolo import yolo
 from zeta.nn.modules.swiglu import SwiGLU, SwiGLUStacked
+from zeta.nn.modules.img_patch_embed import ImgPatchEmbed
 
 # from zeta.nn.modules.img_reshape import image_reshape
 # from zeta.nn.modules.flatten_features import flatten_features
@@ -111,4 +112,5 @@
  "AdaptiveLayerNorm",
  "SwiGLU",
  "SwiGLUStacked",
+ "ImgPatchEmbed",
 ]

zeta/nn/modules/img_patch_embed.py

@@ -0,0 +1,45 @@
+from torch import nn
+
+
+class ImgPatchEmbed(nn.Module):
+ """patch embedding module
+
+
+ Args:
+ img_size (int, optional): image size. Defaults to 224.
+ patch_size (int, optional): patch size. Defaults to 16.
+ in_chans (int, optional): input channels. Defaults to 3.
+ embed_dim (int, optional): embedding dimension. Defaults to 768.
+
+ Examples:
+ >>> x = torch.randn(1, 3, 224, 224)
+ >>> model = ImgPatchEmbed()
+ >>> model(x).shape
+ torch.Size([1, 196, 768])
+
+
+ """
+
+ def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+ super().__init__()
+ num_patches = (img_size // patch_size) * (img_size // patch_size)
+ self.img_size = img_size
+ self.patch_size = patch_size
+ self.num_patches = num_patches
+
+ self.proj = nn.Conv2d(
+ in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
+ )
+
+ def forward(self, x):
+ """Forward
+
+ Args:
+ x (_type_): _description_
+
+ Returns:
+ _type_: _description_
+ """
+ B, C, H, W = x.shape
+ x = self.proj(x).flatten(2).transpose(1, 2)
+ return x

zeta/nn/modules/simple_mamba.py

@@ -6,7 +6,6 @@
 from typing import Optional, Union
 
 
-
 # [HELPERS] ----------------------------------------------------------------------------------------
 class RMSNorm(nn.Module):
  def __init__(self, dim: int, eps: float = 1e-5):
@@ -57,8 +56,6 @@ def forward(self, x):
  return output
 
 
-
-
 class Mamba(nn.Module):
  def __init__(
  self, vocab_size: int = None, dim: int = None, depth: int = None
@@ -98,7 +95,6 @@ class MambaLMHeadModel, https://github.com/state-spaces/mamba/blob/main/mamba_ss
  return logits
 
 
-
 class MambaBlock(nn.Module):
  def __init__(
  self,
@@ -107,7 +103,7 @@ def __init__(
  depth: int,
  d_state: int = 16,
  expand: int = 2,
- dt_rank: Union[int, str] = 'auto',
+ dt_rank: Union[int, str] = "auto",
  d_conv: int = 4,
  conv_bias: bool = True,
  bias: bool = False,
@@ -136,7 +132,6 @@ def __init__(
  self.A_log = nn.Parameter(torch.log(A))
  self.D = nn.Parameter(torch.ones(dim_inner))
  self.out_proj = nn.Linear(dim_inner, dim, bias=bias)
-
 
  def forward(self, x):
  """Mamba block forward. This looks the same as Figure 3 in Section 3.4 in the Mamba paper [1].
@@ -260,4 +255,3 @@ def selective_scan(self, u, delta, A, B, C, D):
  y = y + u * rearrange(D, "d_in -> d_in 1")
 
  return y
-

zeta/utils/disable_logging.py

@@ -1,13 +1,55 @@
+# import logging
+# import os
+# import warnings
+
+
+# def disable_warnings_and_logs():
+# """
+# Disables various warnings and logs.
+# """
+# # disable warnings
+# warnings.filterwarnings("ignore")
+
+# # disable tensorflow warnings
+# os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
+
+# # disable bnb warnings and others
+# logging.getLogger().setLevel(logging.WARNING)
+
+# class CustomFilter(logging.Filter):
+# def filter(self, record):
+# unwanted_logs = [
+# "Setting ds_accelerator to mps (auto detect)",
+# (
+# "NOTE: Redirects are currently not supported in Windows or"
+# " MacOs."
+# ),
+# ]
+# return not any(log in record.getMessage() for log in unwanted_logs)
+
+# # add custom filter to root logger
+# logger = logging.getLogger()
+# f = CustomFilter()
+# logger.addFilter(f)
+
+# # disable specific loggers
+# loggers = [
+# "real_accelerator",
+# "torch.distributed.elastic.multiprocessing.redirects",
+# ]
+
+# for logger_name in loggers:
+# logger = logging.getLogger(logger_name)
+# logger.setLevel(logging.CRITICAL)
+
+
 import logging
 import os
 import warnings
 
-
 def disable_warnings_and_logs():
- """Disable warnings and logs.
-
- Returns:
- _type_: _description_
+ """
+ Disables various warnings and logs.
  """
  # disable warnings
  warnings.filterwarnings("ignore")
@@ -20,12 +62,19 @@ def disable_warnings_and_logs():
 
  class CustomFilter(logging.Filter):
  def filter(self, record):
- msg = "Created a temporary directory at"
- return msg not in record.getMessage()
+ unwanted_logs = [
+ "Setting ds_accelerator to mps (auto detect)",
+ (
+ "NOTE: Redirects are currently not supported in Windows or"
+ " MacOs."
+ ),
+ ]
+ return not any(log in record.getMessage() for log in unwanted_logs)
 
+ # add custom filter to root logger
  logger = logging.getLogger()
  f = CustomFilter()
  logger.addFilter(f)
 
-
-disable_warnings_and_logs()
+ # disable all loggers
+ logging.disable(logging.CRITICAL)