allow for turning off typechecking for torch compile

Author: lucidrains

README.md

@@ -132,6 +132,14 @@ mesh_token_ids = autoencoder.tokenize(
 # (batch, num face vertices, residual quantized layer)
 ```
 
+## Typecheck
+
+At the project root, run
+
+```bash
+$ cp .env.sample .env
+```
+
 ## Todo
 
 - [x] autoencoder

meshgpt_pytorch/data.py

@@ -2,20 +2,19 @@
 
 from pathlib import Path
 from functools import partial
+
 import torch
 from torch import Tensor
 from torch import is_tensor
-import torch.nn.functional as F
 from torch.utils.data import Dataset
 from torch.nn.utils.rnn import pad_sequence
 
-import numpy as np
 from numpy.lib.format import open_memmap
 
 from einops import rearrange, reduce
 
 from beartype import beartype
-from beartype.typing import Tuple, List, Callable, Dict, Callable
+from beartype.typing import Tuple, List, Callable, Dict
 
 from torchtyping import TensorType
 

meshgpt_pytorch/meshgpt_pytorch.py

@@ -11,12 +11,10 @@
 from torch.utils.checkpoint import checkpoint
 from torch.cuda.amp import autocast
 
-from torchtyping import TensorType
-
 from pytorch_custom_utils import save_load
 
-from beartype import beartype
 from beartype.typing import Tuple, Callable, List, Dict, Any
+from meshgpt_pytorch.typing import Float, Int, Bool, typecheck
 
 from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
 
@@ -26,7 +24,6 @@
 from einx import get_at
 
 from x_transformers import Decoder
-from x_transformers.attend import Attend
 from x_transformers.x_transformers import RMSNorm, FeedForward, LayerIntermediates
 
 from x_transformers.autoregressive_wrapper import (
@@ -78,8 +75,8 @@ def divisible_by(num, den):
 def is_odd(n):
     return not divisible_by(n, 2)
 
-def is_empty(l):
-    return len(l) == 0
+def is_empty(x):
+    return len(x) == 0
 
 def is_tensor_empty(t: Tensor):
     return t.numel() == 0
@@ -157,7 +154,7 @@ def derive_angle(x, y, eps = 1e-5):
 
 @torch.no_grad()
 def get_derived_face_features(
-    face_coords: TensorType['b', 'nf', 'nvf', 3, float]  # 3 or 4 vertices with 3 coordinates
+    face_coords: Float['b nf nvf 3']  # 3 or 4 vertices with 3 coordinates
 ):
     shifted_face_coords = torch.cat((face_coords[:, :, -1:], face_coords[:, :, :-1]), dim = 2)
 
@@ -178,7 +175,7 @@ def get_derived_face_features(
 
 # tensor helper functions
 
-@beartype
+@typecheck
 def discretize(
     t: Tensor,
     *,
@@ -194,7 +191,7 @@ def discretize(
 
     return t.round().long().clamp(min = 0, max = num_discrete - 1)
 
-@beartype
+@typecheck
 def undiscretize(
     t: Tensor,
     *,
@@ -210,7 +207,7 @@ def undiscretize(
     t /= num_discrete
     return t * (hi - lo) + lo
 
-@beartype
+@typecheck
 def gaussian_blur_1d(
     t: Tensor,
     *,
@@ -234,7 +231,7 @@ def gaussian_blur_1d(
     out = F.conv1d(t, kernel, padding = half_width, groups = channels)
     return rearrange(out, 'b c n -> b n c')
 
-@beartype
+@typecheck
 def scatter_mean(
     tgt: Tensor,
     indices: Tensor,
@@ -421,7 +418,7 @@ def forward(
 
 @save_load(version = __version__)
 class MeshAutoencoder(Module):
-    @beartype
+    @typecheck
     def __init__(
         self,
         num_discrete_coors = 128,
@@ -671,15 +668,15 @@ def _from_pretrained(
         model.to(map_location)
         return model
 
-    @beartype
+    @typecheck
     def encode(
         self,
         *,
-        vertices:         TensorType['b', 'nv', 3, float],
-        faces:            TensorType['b', 'nf', 'nvf', int],
-        face_edges:       TensorType['b', 'e', 2, int],
-        face_mask:        TensorType['b', 'nf', bool],
-        face_edges_mask:  TensorType['b', 'e', bool],
+        vertices:         Float['b nv 3'],
+        faces:            Int['b nf nvf'],
+        face_edges:       Int['b e 2'],
+        face_mask:        Bool['b nf'],
+        face_edges_mask:  Bool['b e'],
         return_face_coordinates = False
     ):
         """
@@ -692,7 +689,6 @@ def encode(
         d - embed dim
         """
 
-        batch, num_vertices, num_coors, device = *vertices.shape, vertices.device
         _, num_faces, num_vertices_per_face = faces.shape
 
         assert self.num_vertices_per_face == num_vertices_per_face
@@ -773,18 +769,18 @@ def encode(
 
         return face_embed, discrete_face_coords
 
-    @beartype
+    @typecheck
     def quantize(
         self,
         *,
-        faces: TensorType['b', 'nf', 'nvf', int],
-        face_mask: TensorType['b', 'n', bool],
-        face_embed: TensorType['b', 'nf', 'd', float],
+        faces: Int['b nf nvf'],
+        face_mask: Bool['b n'],
+        face_embed: Float['b nf d'],
         pad_id = None,
         rvq_sample_codebook_temp = 1.
     ):
         pad_id = default(pad_id, self.pad_id)
-        batch, num_faces, device = *faces.shape[:2], faces.device
+        batch, device = faces.shape[0], faces.device
 
         max_vertex_index = faces.amax()
         num_vertices = int(max_vertex_index.item() + 1)
@@ -858,11 +854,11 @@ def quantize_wrapper_fn(inp):
 
         return face_embed_output, codes_output, commit_loss
 
-    @beartype
+    @typecheck
     def decode(
         self,
-        quantized: TensorType['b', 'n', 'd', float],
-        face_mask:  TensorType['b', 'n', bool]
+        quantized: Float['b n d'],
+        face_mask:  Bool['b n']
     ):
         conv_face_mask = rearrange(face_mask, 'b n -> b 1 n')
 
@@ -884,12 +880,12 @@ def decode(
 
         return rearrange(x, 'b d n -> b n d')
 
-    @beartype
+    @typecheck
     @torch.no_grad()
     def decode_from_codes_to_faces(
         self,
         codes: Tensor,
-        face_mask: TensorType['b', 'n', bool] | None = None,
+        face_mask: Bool['b n'] | None = None,
         return_discrete_codes = False
     ):
         codes = rearrange(codes, 'b ... -> b (...)')
@@ -964,13 +960,13 @@ def tokenize(self, vertices, faces, face_edges = None, **kwargs):
 
         return codes
 
-    @beartype
+    @typecheck
     def forward(
         self,
         *,
-        vertices:       TensorType['b', 'nv', 3, float],
-        faces:          TensorType['b', 'nf', 'nvf', int],
-        face_edges:     TensorType['b', 'e', 2, int] | None = None,
+        vertices:       Float['b nv 3'],
+        faces:          Int['b nf nvf'],
+        face_edges:     Int['b e 2'] | None = None,
         return_codes = False,
         return_loss_breakdown = False,
         return_recon_faces = False,
@@ -980,7 +976,7 @@ def forward(
         if not exists(face_edges):
             face_edges = derive_face_edges_from_faces(faces, pad_id = self.pad_id)
 
-        num_faces, num_face_edges, device = faces.shape[1], face_edges.shape[1], faces.device
+        device = faces.device
 
         face_mask = reduce(faces != self.pad_id, 'b nf c -> b nf', 'all')
         face_edges_mask = reduce(face_edges != self.pad_id, 'b e ij -> b e', 'all')
@@ -1079,7 +1075,7 @@ def forward(
 
 @save_load(version = __version__)
 class MeshTransformer(Module,PyTorchModelHubMixin):
-    @beartype
+    @typecheck
     def __init__(
         self,
         autoencoder: MeshAutoencoder,
@@ -1270,7 +1266,7 @@ def _from_pretrained(
     def device(self):
         return next(self.parameters()).device
 
-    @beartype
+    @typecheck
     @torch.no_grad()
     def embed_texts(self, texts: str | List[str]):
         single_text = not isinstance(texts, list)
@@ -1287,7 +1283,7 @@ def embed_texts(self, texts: str | List[str]):
 
     @eval_decorator
     @torch.no_grad()
-    @beartype
+    @typecheck
     def generate(
         self,
         prompt: Tensor | None = None,
@@ -1406,9 +1402,9 @@ def generate(
     def forward(
         self,
         *,
-        vertices:       TensorType['b', 'nv', 3, int],
-        faces:          TensorType['b', 'nf', 'nvf', int],
-        face_edges:     TensorType['b', 'e', 2, int] | None = None,
+        vertices:       Int['b nv 3'],
+        faces:          Int['b nf nvf'],
+        face_edges:     Int['b e 2'] | None = None,
         codes:          Tensor | None = None,
         cache:          LayerIntermediates | None = None,
         **kwargs

meshgpt_pytorch/trainer.py

@@ -1,12 +1,10 @@
 from pathlib import Path
 from functools import partial
 from packaging import version
-from contextlib import nullcontext, contextmanager
+from contextlib import nullcontext
 
 import torch
-from torch import nn, Tensor
 from torch.nn import Module
-import torch.nn.functional as F
 from torch.utils.data import Dataset, DataLoader
 from torch.optim.lr_scheduler import _LRScheduler
 
@@ -19,9 +17,8 @@
 from accelerate import Accelerator
 from accelerate.utils import DistributedDataParallelKwargs
 
-from beartype import beartype
-from beartype.door import is_bearable
 from beartype.typing import Tuple, Type, List
+from meshgpt_pytorch.typing import typecheck, beartype_isinstance
 
 from ema_pytorch import EMA
 
@@ -67,7 +64,7 @@ def maybe_del(d: dict, *keys):
 
 @add_wandb_tracker_contextmanager()
 class MeshAutoencoderTrainer(Module):
-    @beartype
+    @typecheck
     def __init__(
         self,
         model: MeshAutoencoder,
@@ -81,7 +78,9 @@ def __init__(
         learning_rate: float = 1e-4,
         weight_decay: float = 0.,
         max_grad_norm: float | None = None,
-        ema_kwargs: dict = dict(),
+        ema_kwargs: dict = dict(
+            use_foreach = True
+        ),
         scheduler: Type[_LRScheduler] | None = None,
         scheduler_kwargs: dict = dict(),
         accelerator_kwargs: dict = dict(),
@@ -147,7 +146,7 @@ def __init__(
             )
 
         if hasattr(dataset, 'data_kwargs') and exists(dataset.data_kwargs):
-            assert is_bearable(dataset.data_kwargs, List[str])
+            assert beartype_isinstance(dataset.data_kwargs, List[str])
             self.data_kwargs = dataset.data_kwargs
         else:
             self.data_kwargs = data_kwargs
@@ -324,7 +323,7 @@ def forward(self):
 
 @add_wandb_tracker_contextmanager()
 class MeshTransformerTrainer(Module):
-    @beartype
+    @typecheck
     def __init__(
         self,
         model: MeshTransformer,
@@ -407,7 +406,7 @@ def __init__(
             )
 
         if hasattr(dataset, 'data_kwargs') and exists(dataset.data_kwargs):
-            assert is_bearable(dataset.data_kwargs, List[str])
+            assert beartype_isinstance(dataset.data_kwargs, List[str])
             self.data_kwargs = dataset.data_kwargs
         else:
             self.data_kwargs = data_kwargs