Upload model.ipynb

train/model.ipynb

@@ -0,0 +1,667 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "id": "dd07a8e6-5809-4bb7-ba3a-bd6c15b22ff2",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/home/user/conda/envs/senv/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+ " from .autonotebook import tqdm as notebook_tqdm\n"
+ ]
+ }
+ ],
+ "source": [
+ "import random\n",
+ "from statistics import mean\n",
+ "from datetime import datetime\n",
+ "from typing import List, Tuple\n",
+ "import copy\n",
+ "\n",
+ "import torch as th\n",
+ "import pytorch_lightning as pl\n",
+ "from pytorch_lightning.callbacks import ModelCheckpoint\n",
+ "from jaxtyping import Float, Float16, Int\n",
+ "\n",
+ "import trimesh as tm\n",
+ "import numpy as np\n",
+ "import numba\n",
+ "\n",
+ "from torch_geometric.nn.conv import GATv2Conv\n",
+ "\n",
+ "import h5py\n",
+ "\n",
+ "# Clone SAP from original repo https://github.com/autonomousvision/shape_as_points.git\n",
+ "from SAP.dpsr import DPSR\n",
+ "from SAP.model import PSR2Mesh"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "59c87491-5650-4c59-8d33-5153d29fb1a9",
+ "metadata": {
+ "tags": []
+ },
+ "source": [
+ "# Constants"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "id": "26d62fb9-dae9-406b-ba30-3fec1a43a29a",
+ "metadata": {
+ "tags": []
+ },
+ "outputs": [],
+ "source": [
+ "th.manual_seed(0)\n",
+ "np.random.seed(0)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "9ab9502f-e822-4475-9c90-019ff28f12d0",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "IS_DEBUG = True"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "id": "7095231b-e8ed-4c4d-997f-8f58664e9877",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "BATCH_SIZE = 1 # BS\n",
+ "LR = 0.001\n",
+ "\n",
+ "IN_DIM = 1 \n",
+ "OUT_DIM = 1\n",
+ "LATENT_DIM = 32\n",
+ "\n",
+ "DROPOUT_PROB = 0.1\n",
+ "\n",
+ "PADDING = 1.2 # Scaling\n",
+ "\n",
+ "GRID_SIZE = 128\n",
+ "SIGMA = 5.0"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "id": "27b7a406-cbb0-4a36-be1e-a8d8aa82c702",
+ "metadata": {
+ "tags": []
+ },
+ "outputs": [],
+ "source": [
+ "DATASET = \"Synthetic\"\n",
+ "LOG_IDX = 14\n",
+ "LOG_VISUALS = not IS_DEBUG\n",
+ "\n",
+ "CHECKPOINTS_PATH = \"./checkpoints/\"\n",
+ "\n",
+ "FIELDS_H5_PATH = f\"./Standart_fields/{DATASET}_fields_32_512.h5\"\n",
+ "PATH_ORIG_H5 = f\"./Standart_h5/{DATASET}.h5\"\n",
+ "PATH_NOISY_H5 = f\"./Standart_h5/{DATASET}_noisy.h5\"\n",
+ "MIN_V_NUMBER = 1_000\n",
+ "MAX_V_NUMBER = 100_000"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1690b667-0af4-465a-8e3c-4a29622e9e66",
+ "metadata": {
+ "tags": []
+ },
+ "source": [
+ "# Data Preparation"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "id": "2e774809-1293-4f80-8350-59ae7fc86cbb",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "@numba.njit\n",
+ "def generate_grid_edge_list(gs: int = 128):\n",
+ " grid_edge_list = []\n",
+ "\n",
+ " for k in range(gs):\n",
+ " for j in range(gs):\n",
+ " for i in range(gs):\n",
+ " current_idx = i + gs*j + k*gs*gs\n",
+ " if (i - 1) >= 0:\n",
+ " grid_edge_list.append([current_idx, i-1 + gs*j + k*gs*gs])\n",
+ " if (i + 1) < gs:\n",
+ " grid_edge_list.append([current_idx, i+1 + gs*j + k*gs*gs])\n",
+ " if (j - 1) >= 0:\n",
+ " grid_edge_list.append([current_idx, i + gs*(j-1) + k*gs*gs])\n",
+ " if (j + 1) < gs:\n",
+ " grid_edge_list.append([current_idx, i + gs*(j+1) + k*gs*gs])\n",
+ " if (k - 1) >= 0:\n",
+ " grid_edge_list.append([current_idx, i + gs*j + (k-1)*gs*gs])\n",
+ " if (k + 1) < gs:\n",
+ " grid_edge_list.append([current_idx, i + gs*j + (k+1)*gs*gs])\n",
+ " return grid_edge_list\n",
+ "\n",
+ "GRID_EDGE_LIST = None"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 7,
+ "id": "4486968b-3416-41c5-9ecd-429f7cf193de",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "class StandartH5DataSet(th.utils.data.Dataset):\n",
+ " \n",
+ " def _load_data(self, key: str):\n",
+ " key_orig = key.replace(\"_n1\", \"\")\n",
+ " key_orig = key_orig.replace(\"_n2\", \"\")\n",
+ " key_orig = key_orig.replace(\"_n3\", \"\")\n",
+ " key_orig = key_orig.replace(\"_noisy\", \"\")\n",
+ "\n",
+ " vertices = th.tensor(self._noisy_meshes_h5[key][\"vertices\"][:], dtype=th.float)\n",
+ " vertices_normals = th.tensor(self._noisy_meshes_h5[key][\"vertices_normals\"][:], dtype=th.float)\n",
+ " vertices_gt = th.tensor(self._orig_meshes_h5[key_orig][\"vertices\"][:], dtype=th.float)\n",
+ " vertices_normals_gt = th.tensor(self._orig_meshes_h5[key_orig][\"vertices_normals\"][:], dtype=th.float)\n",
+ " field_gt = self.dpsr(vertices_gt.unsqueeze(0), vertices_normals_gt.unsqueeze(0)).squeeze(0)\n",
+ "\n",
+ " adj = np.array(self._noisy_meshes_h5[key][\"edge_index\"][:], dtype=np.int64)\n",
+ " adj = th.tensor(adj, dtype=th.int64)\n",
+ " \n",
+ " return vertices, vertices_normals, vertices_gt, vertices_normals_gt, field_gt, adj\n",
+ " \n",
+ " def __init__(self, \n",
+ " orig_meshes_h5: h5py.Group,\n",
+ " noisy_meshes_h5: h5py.Group,\n",
+ " fields_grid_size: int,\n",
+ " min_verts: int,\n",
+ " max_verts: int) -> None:\n",
+ " super().__init__()\n",
+ " \n",
+ " self.dpsr = DPSR([GRID_SIZE, GRID_SIZE, GRID_SIZE], sig=SIGMA)\n",
+ " \n",
+ " self._orig_meshes_h5 = orig_meshes_h5\n",
+ " self._noisy_meshes_h5 = noisy_meshes_h5\n",
+ " \n",
+ " self._fields_grid_size = str(fields_grid_size)\n",
+ " self._min_verts = min_verts\n",
+ " self._max_verts = max_verts\n",
+ " \n",
+ " self._data = {}\n",
+ " self._keys = []\n",
+ " \n",
+ " # filter keys to load only meshes with requested amount of vertices\n",
+ " for key in self._noisy_meshes_h5.keys():\n",
+ " v_number = self._noisy_meshes_h5[key][\"vertices\"].shape[0]\n",
+ " if (v_number >= self._min_verts) and (v_number <= self._max_verts):\n",
+ " self._keys.append(key)\n",
+ " self._keys = np.array(self._keys, dtype=str)\n",
+ " self._loaded = np.full(shape=self._keys.shape, fill_value=False, dtype=bool)\n",
+ " \n",
+ " def __len__(self) -> int:\n",
+ " return self._keys.shape[0]\n",
+ " \n",
+ " def __getitem__(self, index: int) -> Tuple[Float[th.Tensor, \"N 3\"],\n",
+ " Float[th.Tensor, \"N 3\"],\n",
+ " Float[th.Tensor, \"N 3\"],\n",
+ " Float[th.Tensor, \"N 3\"],\n",
+ " Float[th.Tensor, \"GR GR GR\"],\n",
+ " Float[th.Tensor, \"2 E\"]]:\n",
+ " if self._loaded[index] == False:\n",
+ " data = self._load_data(self._keys[index])\n",
+ " self._data[index] = data\n",
+ " self._loaded[index] = True\n",
+ " return copy.deepcopy(self._data[index])\n",
+ " \n",
+ " @property\n",
+ " def fields_grid_size(self):\n",
+ " return int(self._fields_grid_size)\n",
+ " \n",
+ " def renew_grid_size(self, new_grid_size: int):\n",
+ " self._fields_grid_size = str(new_grid_size)\n",
+ " self._loaded = np.full(shape=self._keys.shape, fill_value=False, dtype=bool)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "13c69a49-5107-4d3e-9b14-1d456768f128",
+ "metadata": {
+ "tags": []
+ },
+ "source": [
+ "# Model"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1d9a9aac-d229-489a-844d-a1d1cbd34c56",
+ "metadata": {
+ "tags": []
+ },
+ "source": [
+ "### Form Optimizer "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 8,
+ "id": "940babdc-3e4f-4310-8bfd-48b23d0758dc",
+ "metadata": {
+ "tags": []
+ },
+ "outputs": [],
+ "source": [
+ "class FormOptimizer(th.nn.Module):\n",
+ " def __init__(self) -> None:\n",
+ " super().__init__()\n",
+ " \n",
+ " layers = []\n",
+ " \n",
+ " self.gconv1 = GATv2Conv(in_channels=IN_DIM, out_channels=LATENT_DIM, heads=1, dropout=DROPOUT_PROB)\n",
+ " self.gconv2 = GATv2Conv(in_channels=LATENT_DIM, out_channels=LATENT_DIM, heads=1, dropout=DROPOUT_PROB)\n",
+ " \n",
+ " self.actv = th.nn.Sigmoid()\n",
+ " self.head = th.nn.Linear(in_features=LATENT_DIM, out_features=OUT_DIM)\n",
+ "\n",
+ " def forward(self, \n",
+ " field: Float[th.Tensor, \"GS GS GS\"]) -> Float[th.Tensor, \"GS GS GS\"]:\n",
+ " \"\"\"\n",
+ " Args:\n",
+ " field (Tensor [GS, GS, GS]): vertices and normals tensor.\n",
+ " \"\"\"\n",
+ " vertex_features = field.clone()\n",
+ " vertex_features = vertex_features.reshape(GRID_SIZE*GRID_SIZE*GRID_SIZE, IN_DIM)\n",
+ " \n",
+ " vertex_features = self.gconv1(x=vertex_features, edge_index=GRID_EDGE_LIST) \n",
+ " vertex_features = self.gconv2(x=vertex_features, edge_index=GRID_EDGE_LIST) \n",
+ " field_delta = self.head(self.actv(vertex_features))\n",
+ " \n",
+ " field_delta = field_delta.reshape(BATCH_SIZE, GRID_SIZE, GRID_SIZE, GRID_SIZE)\n",
+ " field_delta += field \n",
+ " field_delta = th.clamp(field_delta, min=-0.5, max=0.5)\n",
+ " \n",
+ " return field_delta"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "67b40c5b-ff1b-416d-b892-c544386eaa95",
+ "metadata": {
+ "toc-hr-collapsed": true
+ },
+ "source": [
+ "### Full"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 9,
+ "id": "bce3aa63-9bd7-4ac8-939d-395d63dd3cad",
+ "metadata": {
+ "scrolled": true,
+ "tags": []
+ },
+ "outputs": [],
+ "source": [
+ "class Model(pl.LightningModule):\n",
+ " def __init__(self):\n",
+ " super().__init__()\n",
+ " self.form_optimizer = FormOptimizer()\n",
+ " \n",
+ " self.dpsr = DPSR([GRID_SIZE, GRID_SIZE, GRID_SIZE], sig=SIGMA)\n",
+ " self.field2mesh = PSR2Mesh().apply\n",
+ "\n",
+ " self.metric = th.nn.MSELoss()\n",
+ "\n",
+ " #video logging databases\n",
+ " dateTimeObj = datetime.now()\n",
+ " start_time = dateTimeObj.strftime(\"%d-%b-%Y_%H-%M-%S\")\n",
+ " \n",
+ " if LOG_VISUALS:\n",
+ " self.h5_frame = 0\n",
+ " self.log_points_file = h5py.File(f\"./logs/points_{start_time}\", \"w\")\n",
+ " self.log_normals_file = h5py.File(f\"./logs/normals_{start_time}\", \"w\")\n",
+ " \n",
+ " self.val_losses = []\n",
+ " self.train_losses = []\n",
+ "\n",
+ " def log_h5(self, points, normals):\n",
+ " dset = self.log_points_file.create_dataset(\n",
+ " name=str(self.h5_frame),\n",
+ " shape=points.shape,\n",
+ " dtype=np.float16, \n",
+ " compression=\"gzip\")\n",
+ " dset[:] = points\n",
+ " dset = self.log_normals_file.create_dataset(\n",
+ " name=str(self.h5_frame),\n",
+ " shape=normals.shape,\n",
+ " dtype=np.float16, \n",
+ " compression=\"gzip\")\n",
+ " dset[:] = normals\n",
+ " self.h5_frame += 1\n",
+ " \n",
+ " def forward(self, \n",
+ " v: Float[th.Tensor, \"BS N 3\"],\n",
+ " n: Float[th.Tensor, \"BS N 3\"]) -> Tuple[Float[th.Tensor, \"BS N 3\"], # v - vertices\n",
+ " Int[th.Tensor, \"2 E\"], # f - faces\n",
+ " Float[th.Tensor, \"BS N 3\"], # n - vertices normals\n",
+ " Float[th.Tensor, \"BS GR GR GR\"]]: # field: \n",
+ " field = self.dpsr(v, n)\n",
+ " field = self.form_optimizer(field)\n",
+ " v, f, n = self.field2mesh(field)\n",
+ " return v, f, n, field\n",
+ "\n",
+ " def training_step(self, batch, batch_idx) -> Float[th.Tensor, \"1\"]:\n",
+ " vertices, vertices_normals, vertices_gt, vertices_normals_gt, field_gt, adj = batch\n",
+ " \n",
+ " mask = th.rand((vertices.shape[1], ), device=th.device(\"cuda\")) < (random.random() / 2.0 + 0.5)\n",
+ " vertices = vertices[:, mask]\n",
+ " vertices_normals = vertices_normals[:, mask]\n",
+ " \n",
+ " vr, fr, nr, field_r = model(vertices, vertices_normals)\n",
+ " \n",
+ " loss = self.metric(field_r, field_gt)\n",
+ " if LOG_VISUALS and (LOG_IDX == batch_idx):\n",
+ " self.log_h5(vr.squeeze(0).detach().cpu().numpy(), nr.squeeze(0).detach().cpu().numpy())\n",
+ " train_per_step_loss = loss.item()\n",
+ " self.train_losses.append(train_per_step_loss)\n",
+ " \n",
+ " return loss\n",
+ " \n",
+ " def on_train_epoch_end(self):\n",
+ " mean_train_per_epoch_loss = mean(self.train_losses)\n",
+ " self.log(\"mean_train_per_epoch_loss\", mean_train_per_epoch_loss, on_step=False, on_epoch=True)\n",
+ " self.train_losses = []\n",
+ " \n",
+ " def validation_step(self, batch, batch_idx):\n",
+ " vertices, vertices_normals, vertices_gt, vertices_normals_gt, field_gt, adj = batch\n",
+ " \n",
+ " vr, fr, nr, field_r = model(vertices, vertices_normals)\n",
+ " \n",
+ " loss = self.metric(field_r, field_gt)\n",
+ " val_per_step_loss = loss.item()\n",
+ " self.val_losses.append(val_per_step_loss)\n",
+ " return loss\n",
+ " \n",
+ " def on_validation_epoch_end(self):\n",
+ " mean_val_per_epoch_loss = mean(self.val_losses)\n",
+ " self.log(\"mean_val_per_epoch_loss\", mean_val_per_epoch_loss, on_step=False, on_epoch=True)\n",
+ " self.val_losses = []\n",
+ "\n",
+ " def configure_optimizers(self):\n",
+ " optimizer = th.optim.Adam(self.parameters(), lr=LR)\n",
+ " scheduler = th.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)\n",
+ " \n",
+ " return {\n",
+ " \"optimizer\": optimizer,\n",
+ " \"lr_scheduler\": {\n",
+ " \"scheduler\": scheduler, \n",
+ " \"monitor\": \"mean_val_per_epoch_loss\",\n",
+ " \"interval\": \"epoch\",\n",
+ " \"frequency\": 1,\n",
+ " \"strict\": True,\n",
+ " \"name\": None,\n",
+ " }\n",
+ " }\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1fb2c5a5-43ee-4a4e-be08-0dcfcb6816de",
+ "metadata": {
+ "tags": []
+ },
+ "source": [
+ "# Loop"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 10,
+ "id": "c94c6a68-3986-48af-9da5-cab8c02a8b7b",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "checkpoint_callback = ModelCheckpoint(\n",
+ " monitor='mean_val_per_epoch_loss', # monitor the validation loss\n",
+ " mode='min', # mode 'min' to save the lowest monitored value\n",
+ " save_top_k=1, # save only the best checkpoint (top 1)\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 11,
+ "id": "03cdddbc-223e-4d40-9fb0-e663beddefda",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/home/user/conda/envs/senv/lib/python3.9/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /opt/conda/conda-bld/pytorch_1670525551200/work/aten/src/ATen/native/TensorShape.cpp:3190.)\n",
+ " return _VF.meshgrid(tensors, **kwargs) # type: ignore[attr-defined]\n",
+ "/home/user/conda/envs/senv/lib/python3.9/site-packages/lightning_fabric/connector.py:554: UserWarning: 16 is supported for historical reasons but its usage is discouraged. Please set your precision to 16-mixed instead!\n",
+ " rank_zero_warn(\n",
+ "Using 16bit Automatic Mixed Precision (AMP)\n",
+ "GPU available: True (cuda), used: True\n",
+ "TPU available: False, using: 0 TPU cores\n",
+ "IPU available: False, using: 0 IPUs\n",
+ "HPU available: False, using: 0 HPUs\n",
+ "Running in `fast_dev_run` mode: will run the requested loop using 300 batch(es). Logging and checkpointing is suppressed.\n",
+ "You are using a CUDA device ('A100-PCIE-40GB') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
+ "/home/user/conda/envs/senv/lib/python3.9/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:617: UserWarning: Checkpoint directory /home/jovyan/Mashurov/GINSAP/checkpoints exists and is not empty.\n",
+ " rank_zero_warn(f\"Checkpoint directory {dirpath} exists and is not empty.\")\n",
+ "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
+ "\n",
+ " | Name | Type | Params\n",
+ "-------------------------------------------------\n",
+ "0 | form_optimizer | FormOptimizer | 2.4 K \n",
+ "1 | dpsr | DPSR | 0 \n",
+ "2 | metric | MSELoss | 0 \n",
+ "-------------------------------------------------\n",
+ "2.4 K Trainable params\n",
+ "0 Non-trainable params\n",
+ "2.4 K Total params\n",
+ "0.010 Total estimated model params size (MB)\n",
+ "/home/user/conda/envs/senv/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:442: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 64 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
+ " rank_zero_warn(\n",
+ "/home/user/conda/envs/senv/lib/python3.9/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:442: PossibleUserWarning: The dataloader, val_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 64 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
+ " rank_zero_warn(\n"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Epoch 0: 100%|██████████| 60/60 [00:17<00:00, 3.52it/s]\n",
+ "Validation: 0it [00:00, ?it/s]\u001b[A\n",
+ "Validation: 0%| | 0/84 [00:00<?, ?it/s]\u001b[A\n",
+ "Validation DataLoader 0: 0%| | 0/84 [00:00<?, ?it/s]\u001b[A\n",
+ "Validation DataLoader 0: 1%| | 1/84 [00:00<00:09, 8.99it/s]\u001b[A\n",
+ "Validation DataLoader 0: 2%|▏ | 2/84 [00:00<00:10, 7.88it/s]\u001b[A\n",
+ "Validation DataLoader 0: 4%|▎ | 3/84 [00:00<00:10, 7.47it/s]\u001b[A\n",
+ "Validation DataLoader 0: 5%|▍ | 4/84 [00:00<00:11, 7.22it/s]\u001b[A\n",
+ "Validation DataLoader 0: 6%|▌ | 5/84 [00:00<00:11, 7.13it/s]\u001b[A\n",
+ "Validation DataLoader 0: 7%|▋ | 6/84 [00:00<00:10, 7.09it/s]\u001b[A\n",
+ "Validation DataLoader 0: 8%|▊ | 7/84 [00:01<00:11, 6.95it/s]\u001b[A\n",
+ "Validation DataLoader 0: 10%|▉ | 8/84 [00:01<00:11, 6.86it/s]\u001b[A\n",
+ "Validation DataLoader 0: 11%|█ | 9/84 [00:01<00:11, 6.81it/s]\u001b[A\n",
+ "Validation DataLoader 0: 12%|█▏ | 10/84 [00:01<00:11, 6.71it/s]\u001b[A\n",
+ "Validation DataLoader 0: 13%|█▎ | 11/84 [00:01<00:10, 6.64it/s]\u001b[A\n",
+ "Validation DataLoader 0: 14%|█▍ | 12/84 [00:01<00:10, 6.60it/s]\u001b[A\n",
+ "Validation DataLoader 0: 15%|█▌ | 13/84 [00:01<00:10, 6.53it/s]\u001b[A\n",
+ "Validation DataLoader 0: 17%|█▋ | 14/84 [00:02<00:10, 6.48it/s]\u001b[A\n",
+ "Validation DataLoader 0: 18%|█▊ | 15/84 [00:02<00:10, 6.44it/s]\u001b[A\n",
+ "Validation DataLoader 0: 19%|█▉ | 16/84 [00:02<00:10, 6.34it/s]\u001b[A\n",
+ "Validation DataLoader 0: 20%|██ | 17/84 [00:02<00:10, 6.27it/s]\u001b[A\n",
+ "Validation DataLoader 0: 21%|██▏ | 18/84 [00:02<00:10, 6.20it/s]\u001b[A\n",
+ "Validation DataLoader 0: 23%|██▎ | 19/84 [00:03<00:10, 6.16it/s]\u001b[A\n",
+ "Validation DataLoader 0: 24%|██▍ | 20/84 [00:03<00:10, 6.16it/s]\u001b[A\n",
+ "Validation DataLoader 0: 25%|██▌ | 21/84 [00:03<00:10, 6.16it/s]\u001b[A\n",
+ "Validation DataLoader 0: 26%|██▌ | 22/84 [00:03<00:10, 6.16it/s]\u001b[A\n",
+ "Validation DataLoader 0: 27%|██▋ | 23/84 [00:03<00:09, 6.17it/s]\u001b[A\n",
+ "Validation DataLoader 0: 29%|██▊ | 24/84 [00:03<00:09, 6.18it/s]\u001b[A\n",
+ "Validation DataLoader 0: 30%|██▉ | 25/84 [00:04<00:09, 6.20it/s]\u001b[A\n",
+ "Validation DataLoader 0: 31%|███ | 26/84 [00:04<00:09, 6.22it/s]\u001b[A\n",
+ "Validation DataLoader 0: 32%|███▏ | 27/84 [00:04<00:09, 6.25it/s]\u001b[A\n",
+ "Validation DataLoader 0: 33%|███▎ | 28/84 [00:04<00:08, 6.22it/s]\u001b[A\n",
+ "Validation DataLoader 0: 35%|███▍ | 29/84 [00:04<00:08, 6.20it/s]\u001b[A\n",
+ "Validation DataLoader 0: 36%|███▌ | 30/84 [00:04<00:08, 6.19it/s]\u001b[A\n",
+ "Validation DataLoader 0: 37%|███▋ | 31/84 [00:04<00:08, 6.21it/s]\u001b[A\n",
+ "Validation DataLoader 0: 38%|███▊ | 32/84 [00:05<00:08, 6.23it/s]\u001b[A\n",
+ "Validation DataLoader 0: 39%|███▉ | 33/84 [00:05<00:08, 6.25it/s]\u001b[A\n",
+ "Validation DataLoader 0: 40%|████ | 34/84 [00:05<00:08, 6.17it/s]\u001b[A\n",
+ "Validation DataLoader 0: 42%|████▏ | 35/84 [00:05<00:08, 6.12it/s]\u001b[A\n",
+ "Validation DataLoader 0: 43%|████▎ | 36/84 [00:05<00:07, 6.07it/s]\u001b[A\n",
+ "Validation DataLoader 0: 44%|████▍ | 37/84 [00:06<00:07, 6.08it/s]\u001b[A\n",
+ "Validation DataLoader 0: 45%|████▌ | 38/84 [00:06<00:07, 6.10it/s]\u001b[A\n",
+ "Validation DataLoader 0: 46%|████▋ | 39/84 [00:06<00:07, 6.11it/s]\u001b[A\n",
+ "Validation DataLoader 0: 48%|████▊ | 40/84 [00:06<00:07, 6.11it/s]\u001b[A\n",
+ "Validation DataLoader 0: 49%|████▉ | 41/84 [00:06<00:07, 6.11it/s]\u001b[A\n",
+ "Validation DataLoader 0: 50%|█████ | 42/84 [00:06<00:06, 6.11it/s]\u001b[A\n",
+ "Validation DataLoader 0: 51%|█████ | 43/84 [00:07<00:06, 6.11it/s]\u001b[A\n",
+ "Validation DataLoader 0: 52%|█████▏ | 44/84 [00:07<00:06, 6.10it/s]\u001b[A\n",
+ "Validation DataLoader 0: 54%|█████▎ | 45/84 [00:07<00:06, 6.07it/s]\u001b[A\n",
+ "Validation DataLoader 0: 55%|█████▍ | 46/84 [00:07<00:06, 6.06it/s]\u001b[A\n",
+ "Validation DataLoader 0: 56%|█████▌ | 47/84 [00:07<00:06, 6.05it/s]\u001b[A\n",
+ "Validation DataLoader 0: 57%|█████▋ | 48/84 [00:07<00:05, 6.06it/s]\u001b[A\n",
+ "Validation DataLoader 0: 58%|█████▊ | 49/84 [00:08<00:05, 6.04it/s]\u001b[A\n",
+ "Validation DataLoader 0: 60%|█████▉ | 50/84 [00:08<00:05, 6.02it/s]\u001b[A\n",
+ "Validation DataLoader 0: 61%|██████ | 51/84 [00:08<00:05, 6.01it/s]\u001b[A\n",
+ "Validation DataLoader 0: 62%|██████▏ | 52/84 [00:08<00:05, 6.03it/s]\u001b[A\n",
+ "Validation DataLoader 0: 63%|██████▎ | 53/84 [00:08<00:05, 6.05it/s]\u001b[A\n",
+ "Validation DataLoader 0: 64%|██████▍ | 54/84 [00:08<00:04, 6.07it/s]\u001b[A\n",
+ "Validation DataLoader 0: 65%|██████▌ | 55/84 [00:09<00:04, 6.07it/s]\u001b[A\n",
+ "Validation DataLoader 0: 67%|██████▋ | 56/84 [00:09<00:04, 6.07it/s]\u001b[A\n",
+ "Validation DataLoader 0: 68%|██████▊ | 57/84 [00:09<00:04, 6.07it/s]\u001b[A\n",
+ "Validation DataLoader 0: 69%|██████▉ | 58/84 [00:09<00:04, 6.06it/s]\u001b[A\n",
+ "Validation DataLoader 0: 70%|███████ | 59/84 [00:09<00:04, 6.05it/s]\u001b[A\n",
+ "Validation DataLoader 0: 71%|███████▏ | 60/84 [00:09<00:03, 6.04it/s]\u001b[A\n",
+ "Validation DataLoader 0: 73%|███████▎ | 61/84 [00:10<00:03, 6.02it/s]\u001b[A\n",
+ "Validation DataLoader 0: 74%|███████▍ | 62/84 [00:10<00:03, 6.00it/s]\u001b[A\n",
+ "Validation DataLoader 0: 75%|███████▌ | 63/84 [00:10<00:03, 5.99it/s]\u001b[A\n",
+ "Validation DataLoader 0: 76%|███████▌ | 64/84 [00:10<00:03, 5.99it/s]\u001b[A\n",
+ "Validation DataLoader 0: 77%|███████▋ | 65/84 [00:10<00:03, 6.00it/s]\u001b[A\n",
+ "Validation DataLoader 0: 79%|███████▊ | 66/84 [00:10<00:02, 6.00it/s]\u001b[A\n",
+ "Validation DataLoader 0: 80%|███████▉ | 67/84 [00:11<00:02, 6.00it/s]\u001b[A\n",
+ "Validation DataLoader 0: 81%|████████ | 68/84 [00:11<00:02, 6.00it/s]\u001b[A\n",
+ "Validation DataLoader 0: 82%|████████▏ | 69/84 [00:11<00:02, 6.00it/s]\u001b[A\n",
+ "Validation DataLoader 0: 83%|████████▎ | 70/84 [00:11<00:02, 6.01it/s]\u001b[A\n",
+ "Validation DataLoader 0: 85%|████████▍ | 71/84 [00:11<00:02, 6.02it/s]\u001b[A\n",
+ "Validation DataLoader 0: 86%|████████▌ | 72/84 [00:11<00:01, 6.03it/s]\u001b[A\n",
+ "Validation DataLoader 0: 87%|████████▋ | 73/84 [00:12<00:01, 6.04it/s]\u001b[A\n",
+ "Validation DataLoader 0: 88%|████████▊ | 74/84 [00:12<00:01, 6.05it/s]\u001b[A\n",
+ "Validation DataLoader 0: 89%|████████▉ | 75/84 [00:12<00:01, 6.06it/s]\u001b[A\n",
+ "Validation DataLoader 0: 90%|█████████ | 76/84 [00:12<00:01, 6.06it/s]\u001b[A\n",
+ "Validation DataLoader 0: 92%|█████████▏| 77/84 [00:12<00:01, 6.07it/s]\u001b[A\n",
+ "Validation DataLoader 0: 93%|█████████▎| 78/84 [00:12<00:00, 6.08it/s]\u001b[A\n",
+ "Validation DataLoader 0: 94%|█████████▍| 79/84 [00:12<00:00, 6.09it/s]\u001b[A\n",
+ "Validation DataLoader 0: 95%|█████████▌| 80/84 [00:13<00:00, 6.10it/s]\u001b[A\n",
+ "Validation DataLoader 0: 96%|█████████▋| 81/84 [00:13<00:00, 6.11it/s]\u001b[A\n",
+ "Validation DataLoader 0: 98%|█████████▊| 82/84 [00:13<00:00, 6.09it/s]\u001b[A\n",
+ "Validation DataLoader 0: 99%|█████████▉| 83/84 [00:13<00:00, 6.09it/s]\u001b[A\n",
+ "Validation DataLoader 0: 100%|██████████| 84/84 [00:13<00:00, 6.08it/s]\u001b[A\n",
+ "Epoch 0: 100%|██████████| 60/60 [00:30<00:00, 1.94it/s] \u001b[A"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "`Trainer.fit` stopped: `max_epochs=1` reached.\n"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Epoch 0: 100%|██████████| 60/60 [00:30<00:00, 1.94it/s]\n"
+ ]
+ }
+ ],
+ "source": [
+ "if __name__ == \"__main__\":\n",
+ " \n",
+ " GRID_EDGE_LIST = generate_grid_edge_list(GRID_SIZE)\n",
+ " GRID_EDGE_LIST = th.tensor(GRID_EDGE_LIST, dtype=th.int)\n",
+ " GRID_EDGE_LIST = GRID_EDGE_LIST.T\n",
+ " GRID_EDGE_LIST = GRID_EDGE_LIST.to(th.device(\"cuda\"))\n",
+ " \n",
+ " noisy_meshes_h5 = h5py.File(\"./Standart_h5/Synthetic_noisy.h5\", \"r\")\n",
+ " orig_meshes_h5 = h5py.File(\"./Standart_h5/Synthetic.h5\", \"r\")\n",
+ " \n",
+ " train_dataset = StandartH5DataSet(orig_meshes_h5=orig_meshes_h5['train'],\n",
+ " noisy_meshes_h5=noisy_meshes_h5['train'],\n",
+ " fields_grid_size=GRID_SIZE,\n",
+ " min_verts=MIN_V_NUMBER,\n",
+ " max_verts=MAX_V_NUMBER)\n",
+ " test_dataset = StandartH5DataSet(orig_meshes_h5=orig_meshes_h5['test'],\n",
+ " noisy_meshes_h5=noisy_meshes_h5['test'],\n",
+ " fields_grid_size=GRID_SIZE,\n",
+ " min_verts=MIN_V_NUMBER,\n",
+ " max_verts=MAX_V_NUMBER)\n",
+ "\n",
+ " train_dataloader = th.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
+ " test_dataloader = th.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)\n",
+ "\n",
+ " trainer = pl.Trainer(accelerator=\"gpu\", \n",
+ " callbacks=[checkpoint_callback],\n",
+ " log_every_n_steps=len(train_dataset)+len(test_dataset),\n",
+ " fast_dev_run=(300 if IS_DEBUG else False),\n",
+ " max_epochs=200,\n",
+ " precision=16)\n",
+ " \n",
+ " model = Model()\n",
+ " trainer.fit(model, train_dataloaders=train_dataloader, val_dataloaders=test_dataloader)\n",
+ " if LOG_VISUALS:\n",
+ " model.log_points_file.close()\n",
+ " model.log_normals_file.close()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "bda6c1bf-7674-4e59-8cc7-dfcba9d689d9",
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "senv",
+ "language": "python",
+ "name": "senv"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.9.16"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}