prototype pipeline

test.py

@@ -0,0 +1,95 @@
+import os
+import sys
+
+import torch
+
+import torch.optim as optim
+import torch.nn as nn
+
+from torch.utils.data import DataLoader
+from data.dataset import Dataset
+from utils.utils import Accumulator
+from models.depth_estim import DepthEstimation
+
+
+def train_step(batch):
+    ### Training
+    model.train() # train mode is on by default after construction
+
+    # 1. Forward pass
+    y_pred = model(batch['mesh_posed'])
+
+    # 2. Calculate loss
+    loss = loss_fn(y_pred, batch['mesh_posed'])
+
+    # 3. Zero grad optimizer
+    optimizer.zero_grad()
+
+    # 4. Loss backward
+    loss.backward()
+
+    # 5. Step the optimizer
+    optimizer.step()
+
+    return { 'train/loss': loss.detach().cpu().numpy() }
+
+
+def valid_step(batch):
+    ### Testing
+    model.eval() # put the model in evaluation mode for testing (inference)
+    # 1. Forward pass
+    with torch.inference_mode():
+        pred = model(batch['mesh_posed'])
+    
+        # 2. Calculate the loss
+        loss = loss_fn(pred, batch['mesh_posed'])
+
+    return { 'valid/loss': loss.detach().cpu().numpy() }
+
+
+def batch_to_device(batch, device):
+    filter = {'mesh_posed'}
+    for key in batch:
+        if key in filter:
+            batch[key] = batch[key].to(device)
+
+
+if __name__ == '__main__':
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+    train_dataset = Dataset('/home/workspace/local-data/train')
+    valid_dataset = Dataset('/home/workspace/local-data/valid')
+    train_dataloader = DataLoader(train_dataset, batch_size=4, drop_last=True, num_workers=2)
+    valid_dataloader = DataLoader(valid_dataset, batch_size=4, drop_last=True, num_workers=2)
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
+    model = DepthEstimation()
+
+    loss_fn = nn.MSELoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+    iteration = 0
+    while True:
+
+        for batch in train_dataloader:
+            train_loss = train_step(batch)
+            train_accumulator.update(train_loss)
+
+            if iteration % 1000 == 0:
+                print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}")
+                train_accumulator.reset()
+
+            iteration += 1
+
+        for batch in valid_dataloader:
+            valid_loss = valid_step(batch)
+            valid_accumulator.update(valid_loss)
+
+            print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}")
+            valid_accumulator.reset()
+
+
+
+    pass