.

Author: KhoiDOO

config/pointcloud/mesh500/mesh500_2048_vqvae_mhvq.yaml

@@ -0,0 +1,78 @@
+name: mesh500_2048_vqvae_mhvq
+exp_root_dir: outputs
+exp_dir: ""
+trial_dir: ""
+n_gpus: 1
+seed: 42
+
+dataset_name: mesh500_2048
+dataset_source: hf
+dataset_img_key: null
+dataset_kwargs:
+  root: ./.cache/${dataset_name}
+
+model: VQVAE
+model_config: VQVAEConfig
+model_kwargs:
+  dim: 128
+  in_channel: 3
+  out_channel: 3
+  layers: 2
+  layer_mults: null
+  num_res_blocks: 1
+  group: 8
+  conv_type: "conv1d"
+  enc_act_func: "LeakyReLU"
+  dec_act_func: "GLU"
+  enc_act_kwargs: {"negative_slope": 0.1}
+  dec_act_kwargs: {"dim": 1}
+  first_conv_kernel_size: 5
+  quantizer: "VectorQuantize"
+  codebook_size: 512
+  quantizer_kwargs: {
+    "codebook_dim": 64,
+    "heads": 8,
+    "decay" : 0.99,
+    "commitment_weight": 0.25,
+    "kmeans_init": True,
+    "use_cosine_sim": True
+  }
+  l2_recon_loss: True
+
+trainer: PCVQVAETrainer
+trainer_config: PCVQVAETrainerConfig
+trainer_kwargs:
+  num_train_steps: 300000
+  batch_size: 128
+  num_workers: 16
+  pin_memory: True
+  grad_accum_every: 1
+  learning_rate: 0.001
+  weight_decay: 0.
+  max_grad_norm: 0.5
+  val_every: 1000
+  val_num_batches: 20
+  val_num_images: 32
+  scheduler: CosineAnnealingLR
+  scheduler_kwargs:
+    T_max: "${sub: ${trainer_kwargs.num_train_steps}, ${trainer_kwargs.warmup_steps}}"
+    eta_min: 0.0005
+  ema_kwargs: null
+  accelerator_kwargs: {}
+  optimizer_name: Adam
+  optimizer_kwargs: {}
+  loss_lambda:
+    recon_loss: 1.
+    quantizer_loss: 1.
+  checkpoint_every: null
+  warmup_steps: 0
+  use_wandb_tracking: False
+  resume: False
+  from_checkpoint: null
+  from_checkpoint_type: null
+
+wandb:
+  project_name: "vitvqganvae"
+  run_name: null
+  kwargs:
+    entity: "heartbeats"

docker/pointcloud/mesh500/mesh500_2048_vqvae_mhvq.dockerfile

@@ -0,0 +1,21 @@
+FROM kohido/base_dl_cuda129:v0.0.6
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+WORKDIR /vitvqganvae
+
+RUN ls
+
+COPY ./config /vitvqganvae/config
+COPY ./vitvqganvae /vitvqganvae/vitvqganvae
+COPY ./main.py /vitvqganvae/main.py
+
+CMD wandb login ${WANDB_API_KEY} && accelerate launch \
+    --mixed_precision=no \
+    --num_processes=1 \
+    --num_machines=1 \
+    --dynamo_backend=no \
+    main.py \
+    --config config/pointcloud/mesh500/mesh500_2048_vqvae_mhvq.yaml \
+    --train \
+    trainer_kwargs.use_wandb_tracking=True

main.py

@@ -67,7 +67,12 @@ def main(args, extras):
     model_module = model_cls(**model_config)
     try:
         sample: Tensor = train_ds[0].unsqueeze(0)
-        summary(copy.deepcopy(model_module), input_size=sample.shape)
+        summary(
+            copy.deepcopy(model_module),
+            input_data=sample,
+            col_names=["input_size", "output_size", "num_params", "params_percent", "trainable"],
+            # depth=2
+        )
     except Exception as e:
         print(f"Cannot run model summary: {e}")
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vitvqganvae"
-version = "0.4.7"
+version = "0.4.8"
 authors = [
   { name="KhoiDOO", email="khoido8899@gmail.com" },
 ]

vitvqganvae/data/hf/mesh500.py

@@ -13,20 +13,26 @@ class Mesh500(Dataset):
     def __init__(self, root: str, num_points: int = 1024):
         super().__init__()
 
-        if num_points not in [1024, 4096]:
-            raise ValueError("num_points should be one of 1024 or 4096 for Mesh500 dataset")
+        if num_points not in [1024, 2048, 4096]:
+            raise ValueError("num_points should be one of 1024, 2048, or 4096 for Mesh500 dataset")
 
         self._root = root
         self._num_points = num_points
-        self._dataset = load_dataset(f"kohido/mesh500_{num_points}pts", cache_dir=self._root)['train']['points']
+        if self._num_points == 2048:
+            self._dataset = load_dataset(f"kohido/mesh500_4096pts", cache_dir=self._root)['train']['points']
+        elif self._num_points == 4096:
+            self._dataset = load_dataset(f"kohido/mesh500_4096pts", cache_dir=self._root)['train']['points']
 
     def __len__(self) -> int:
         return len(self._dataset)
 
     def __getitem__(self, index: int) -> Tensor:
         points = self._dataset[index]
-        points: np.ndarray = np.array(points) # (1024, 3)
+        points: np.ndarray = np.array(points) # (self._num_points, 3)
         points = points[np.lexsort((points[:, 2], points[:, 1], points[:, 0]))]
+        if self._num_points == 2048:
+            # only take even indices
+            points = points[::2]
         # scale to a [-0.5, 0.5] cube
         points = points - np.mean(points, axis=0, keepdims=True)
         max_abs = np.max(np.abs(points))
@@ -54,6 +60,9 @@ def get_mesh500(root: str | None = None, num_points: int = 1024, split: float =
 def get_mesh500_1024(root: str | None = None, split: float = 0.8) -> tuple[Mesh500, Mesh500]:
     return get_mesh500(root=root, num_points=1024, split=split)
 
+def get_mesh500_2048(root: str | None = None, split: float = 0.8) -> tuple[Mesh500, Mesh500]:
+    return get_mesh500(root=root, num_points=2048, split=split)
+
 def get_mesh500_4096(root: str | None = None, split: float = 0.8) -> tuple[Mesh500, Mesh500]:
     return get_mesh500(root=root, num_points=4096, split=split)