100 lines
2.8 KiB
Python
100 lines
2.8 KiB
Python
import os
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
|
|
|
|
import sys
|
|
|
|
import torch
|
|
|
|
import torch.optim as optim
|
|
import torch.nn as nn
|
|
|
|
from torch.utils.data import DataLoader
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
|
|
from data.dataset import Dataset
|
|
from utils.utils import Accumulator
|
|
from models.depth_estim import DepthEstimation
|
|
from models.convnext import ConvNeXt
|
|
|
|
|
|
def train_step(batch):
|
|
### Training
|
|
model.train() # train mode is on by default after construction
|
|
|
|
# 1. Forward pass
|
|
y_pred = model([batch['rgb_image'], batch['view']])
|
|
|
|
# 2. Calculate loss
|
|
loss = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
|
|
|
|
# 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['rgb_image'], batch['view']])
|
|
|
|
# 2. Calculate the loss
|
|
loss = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
|
|
|
|
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('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt')
|
|
valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt')
|
|
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()
|
|
|
|
model = DepthEstimation(device=device)
|
|
|
|
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 |