basic test

test.py

@@ -17,10 +17,10 @@ def train_step(batch):
     model.train() # train mode is on by default after construction
 
     # 1. Forward pass
-    y_pred = model(batch['mesh_posed'])
+    y_pred = model(batch['rgb_image'])
 
     # 2. Calculate loss
-    loss = loss_fn(y_pred, batch['mesh_posed'])
+    loss = loss_fn(y_pred, batch['depth_image'][:,None,:,:])
 
     # 3. Zero grad optimizer
     optimizer.zero_grad()
@@ -39,10 +39,10 @@ def valid_step(batch):
     model.eval() # put the model in evaluation mode for testing (inference)
     # 1. Forward pass
     with torch.inference_mode():
-        pred = model(batch['mesh_posed'])
+        pred = model(batch['rgb_image'])
     
         # 2. Calculate the loss
-        loss = loss_fn(pred, batch['mesh_posed'])
+        loss = loss_fn(pred, batch['depth_image'][:,None,:,:])
 
     return { 'valid/loss': loss.detach().cpu().numpy() }
 
@@ -59,8 +59,8 @@ if __name__ == '__main__':
 
     train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt')
     valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt')
-    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_dataloader = DataLoader(train_dataset, batch_size=2, drop_last=True, num_workers=0)
+    valid_dataloader = DataLoader(valid_dataset, batch_size=2, drop_last=True, num_workers=0)
     train_accumulator = Accumulator()
     valid_accumulator = Accumulator()
 
@@ -86,8 +86,8 @@ if __name__ == '__main__':
             valid_loss = valid_step(batch)
             valid_accumulator.update(valid_loss)
 
-            print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}")
-            valid_accumulator.reset()
+        print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}")
+        valid_accumulator.reset()