experiment.py

import torch
from torch.optim import Adam
from torch.optim.lr_scheduler import MultiStepLR
from torch.utils.data import DataLoader
import pytorch_lightning as pl

from utils.metrics import SceneFlowMetrics
from utils.utils import get_num_workers
from utils.exp_utils import get_datasets
from losses import *


class SceneFlowExp(pl.LightningModule):
    """
    Class for Scene Flow experiment.
    """

    def __init__(self, model, hparams):
        """
        Initializes an experiment.

        Arguments:
            model : a torch.nn.Module object, model to be tested
            hparams : a dictionary of hyper parameters
        """

        super(SceneFlowExp, self).__init__()
        self.model = model
        self.hparams = hparams
        self.loss = losses_dict[hparams['loss']['loss_type']](**hparams['loss'])
        self.train_metrics = SceneFlowMetrics(split='train', loss_params=hparams['loss'], reduce_op='mean')
        self.val_metrics = SceneFlowMetrics(split='val', loss_params=hparams['loss'], reduce_op='mean')

    def forward(self, pos1, pos2, feat1, feat2, iters):
        """
        A forward call
        """
        return self.model(pos1, pos2, feat1, feat2, iters)

    def sequence_loss(self, pos1, pos2, flows_pred, flow_gt):
        if 'loss_iters_w' in self.hparams:
            assert (len(self.hparams['loss_iters_w']) == len(flows_pred))
            loss = torch.zeros(1).cuda()
            for i, w in enumerate(self.hparams['loss_iters_w']):
                loss += w * self.loss(pos1, pos2, flows_pred[i], flow_gt, i)
        else:
            loss = self.loss(pos1, pos2, flows_pred[-1], flow_gt)
        return loss
        
    def training_step(self, batch, batch_idx):
        """
        Executes a single training step
        """

        pos1, pos2, feat1, feat2, flow_gt, fnames = batch
        flows_pred = self(pos1, pos2, feat1, feat2, self.hparams['train_iters'])
        loss = self.sequence_loss(pos1, pos2, flows_pred, flow_gt)
        metrics = self.train_metrics(pos1, pos2, flows_pred, flow_gt)

        train_results = pl.TrainResult(loss)
        train_results.log('train_loss', loss, sync_dist=True, on_step=False, on_epoch=True, logger=True, prog_bar=False, reduce_fx=torch.mean)
        train_results.log_dict(metrics, on_step=False, on_epoch=True, logger=True, prog_bar=False, reduce_fx=torch.mean)  # No need to sync_dist since metrics are already synced     
        
        return train_results

    def _test_val_step(self, batch, batch_idx, split):
        pos1, pos2, feat1, feat2, flow_gt, fnames = batch
        flows_pred = self(pos1, pos2, feat1, feat2, self.hparams[f'{split}_iters'])
        loss = self.sequence_loss(pos1, pos2, flows_pred, flow_gt)
        metrics = self.val_metrics(pos1, pos2, flows_pred, flow_gt)
        
        i_last = self.hparams[f'{split}_iters'] - 1
        val_results = pl.EvalResult(checkpoint_on=metrics[f'val_epe3d_i#{i_last}'])
        val_results.log('val_loss', loss, sync_dist=True, on_step=False, on_epoch=True, logger=True, prog_bar=False, reduce_fx=torch.mean)
        val_results.log_dict(metrics, on_step=False, on_epoch=True, logger=True, prog_bar=False, reduce_fx=torch.mean)  # No need to sync_dist since metrics are already synced

        return val_results

    def validation_step(self, batch, batch_idx):
        """
        Executes a single validation step.
        """
        return self._test_val_step(batch, batch_idx, 'val')

    def test_step(self, batch, batch_idx):
        """
        Executes a single test step.
        """
        return self._test_val_step(batch, batch_idx, 'test')

    def configure_optimizers(self):
        optimizer = Adam(self.parameters(), lr=self.hparams['optimizer']['lr'])
        lr_scheduler = MultiStepLR(optimizer, milestones=self.hparams['scheduler']['milestones'], gamma=self.hparams['scheduler']['gamma'])

        return {'optimizer': optimizer,
                'lr_scheduler': lr_scheduler}

    def train_dataloader(self):
        """
        Returns a train set Dataloader object 
        """
        return self._dataloader(split_type='train')
        
    def val_dataloader(self):
        """
        Returns a validation set Dataloader object 
        """
        return self._dataloader(split_type='val')

    def test_dataloader(self):
        """
        Returns a validation set Dataloader object
        """
        return self._dataloader(split_type='test')

    # def prepare_data(self):
    #     """
    #     Prepare data, executed before loading to GPUs (globally or per node).
    #     """

    def setup(self, stage: str):
        """
        Load datasets
        """
        train_dataset, val_dataset, test_dataset = get_datasets(self.hparams['data'])
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset
        self.test_dataset = test_dataset

    def _dataloader(self, split_type: str):
        """
        Arguments:
            split_type : (str) shouild be one of ['train', 'val', 'test']
        Return:
            A DataLoader object of the correspondig split
        """
        
        split_dict = {
            'train': self.train_dataset,
            'val': self.val_dataset,
            'test': self.test_dataset
            }
        is_train = (split_type == 'train')
        num_workers = get_num_workers(self.hparams['num_workers'])
        if split_dict[split_type] is None:
            loader = None
        else:
            loader = DataLoader(split_dict[split_type],
                                batch_size=self.hparams['batch_size'],
                                num_workers=num_workers,
                                pin_memory=True,
                                shuffle=True if is_train else False)
        return loader