train_depth

#!/usr/bin/env python

import os
import numpy as np
import torch
from torch.utils.data import DataLoader, BatchSampler
from torch.utils.data.sampler import SequentialSampler
from torch.utils.tensorboard import SummaryWriter
from argparse import ArgumentParser
import datasets
from tqdm import tqdm
from time import time
import segmentation_models_pytorch as smp
from typing import Iterator, List
from traversability_estimation.utils import visualize_imgs, visualize_cloud, create_model
from datasets.base_dataset import VOID_VALUE
import torchmetrics


def parse_arguments():
    parser = ArgumentParser()
    parser.add_argument('--lr', type=float, default=1e-4)
    parser.add_argument('--datasets', nargs='+', type=str, default=['Rellis3DClouds'])
    parser.add_argument('--output', type=str, default=None)
    parser.add_argument('--dont_save_models', action='store_true')
    parser.add_argument('--architecture', type=str, default='deeplabv3_resnet101')
    parser.add_argument('--pretrained_weights', type=str, default=None)
    parser.add_argument('--batch_size', type=int, default=4)
    parser.add_argument('--n_epochs', type=int, default=100)
    parser.add_argument('--n_workers', type=int, default=os.cpu_count() // 2)
    parser.add_argument('--data_fields', nargs='+', type=str, default=['depth'])
    parser.add_argument('--n_samples', type=int, default=None)
    parser.add_argument('--vis_preds', action='store_true')
    parser.add_argument('--loss_fn', type=str, default='lovasz')
    args = parser.parse_args()

    return args


class CommonBatchSampler(BatchSampler):
    """
    Selecting indices from two different datasets to form batches of data
    """

    def __init__(self, sampler, batch_size, drop_last, shuffle=False):
        super(CommonBatchSampler, self).__init__(sampler=sampler, batch_size=batch_size, drop_last=drop_last)
        self.shuffle = shuffle

    def __iter__(self) -> Iterator[List[int]]:
        datasets_idx_border = self.sampler.data_source.cumulative_sizes[0]

        # assume we have 2 datasets two concatenate
        assert len(self.sampler.data_source.cumulative_sizes) == 2
        if self.shuffle:
            # random batch sampler: select batches randomly either from one dataset or another
            ids = {'0': list(self.sampler)[:datasets_idx_border],
                   '1': list(self.sampler)[datasets_idx_border:]}
            while len(ids['0']) > 0 or len(ids['1']) > 0:
                data_id = np.random.choice(list(ids.keys()))
                if len(ids[data_id]) == 0:
                    # inverse 0 to 1 and vise versa
                    data_id = str(int(not int(data_id)))
                assert len(ids[data_id]) > 0

                if len(ids[data_id]) > self.batch_size:
                    batch = np.random.choice(ids[data_id], self.batch_size, replace=False).tolist()
                    assert len(batch) > 0
                else:
                    if self.drop_last:
                        # TODO: in case drop_last = True goes out of the cycle in the middle of dataloader
                        break
                    else:
                        batch = ids[data_id].copy()
                        assert len(batch) > 0

                for ind in batch:
                    ids[data_id].remove(ind)
                assert len(batch) > 0
                yield batch
        else:
            # sequential batch sampler
            batch = []
            for idx in self.sampler:
                batch.append(idx)
                if len(batch) == self.batch_size or idx == datasets_idx_border - 1:
                    yield batch
                    batch = []
            if len(batch) > 0 and not self.drop_last:
                yield batch


def create_dataloaders(args):
    print('Using datasets for training: %s' % ' '.join(args.datasets))

    Dataset = eval('datasets.%s' % args.datasets[0])
    train_dataset = Dataset(split='train',
                            output=args.output,
                            fields=args.data_fields, num_samples=args.n_samples,
                            lidar_W_step=1,
                            labels_mode='labels')
    valid_dataset = Dataset(split='val',
                            output=args.output,
                            fields=args.data_fields, num_samples=args.n_samples,
                            lidar_W_step=1,
                            labels_mode='labels')

    train_datasets, valid_datasets = [train_dataset], [valid_dataset]

    if len(args.datasets) > 1:
        assert len(args.datasets) == 2

        DatasetFT = eval('datasets.%s' % args.datasets[1])

        ft_dataset_train = DatasetFT(split='train',
                                     output=args.output,
                                     fields=args.data_fields, num_samples=args.n_samples,
                                     labels_mode='labels')
        assert ft_dataset_train.output == train_dataset.output

        train_dataset_combined = torch.utils.data.ConcatDataset([train_dataset, ft_dataset_train])

        if train_dataset[0][0].shape == ft_dataset_train[0][0].shape:
            train_loader = DataLoader(train_dataset_combined,
                                      batch_size=args.batch_size,
                                      shuffle=True,
                                      num_workers=args.n_workers,
                                      drop_last=True,
                                      pin_memory=True)
        else:
            print('Source and target datasets have different data shapes. Using BatchSampler.')
            # https://stackoverflow.com/questions/51837110/pytorch-data-loading-from-multiple-different-sized-datasets
            batch_sampler = CommonBatchSampler(SequentialSampler(train_dataset_combined),
                                               batch_size=args.batch_size,
                                               drop_last=True,
                                               shuffle=True)
            train_loader = DataLoader(dataset=train_dataset_combined,
                                      num_workers=args.n_workers,
                                      batch_sampler=batch_sampler,
                                      pin_memory=True)

        ft_dataset_val = DatasetFT(split='val',
                                   output=args.output,
                                   fields=args.data_fields, num_samples=args.n_samples,
                                   labels_mode='labels')
        assert ft_dataset_val.output == valid_dataset.output

        valid_dataset_combined = torch.utils.data.ConcatDataset([valid_dataset, ft_dataset_val])

        if valid_dataset[0][0].shape == ft_dataset_val[0][0].shape:
            valid_loader = DataLoader(valid_dataset_combined,
                                      batch_size=1,
                                      shuffle=False,
                                      num_workers=args.n_workers,
                                      pin_memory=True)
        else:
            # https://stackoverflow.com/questions/51837110/pytorch-data-loading-from-multiple-different-sized-datasets
            batch_sampler_val = CommonBatchSampler(SequentialSampler(valid_dataset_combined),
                                                   batch_size=args.batch_size,
                                                   drop_last=False,
                                                   shuffle=False)
            valid_loader = DataLoader(dataset=valid_dataset_combined,
                                      num_workers=args.n_workers,
                                      batch_sampler=batch_sampler_val,
                                      pin_memory=True)

        train_datasets.append(ft_dataset_train)
        valid_datasets.append(ft_dataset_val)
    else:
        assert len(args.datasets) == 1

        train_loader = DataLoader(train_dataset,
                                  batch_size=args.batch_size,
                                  shuffle=True,
                                  num_workers=args.n_workers,
                                  drop_last=True,
                                  pin_memory=True)

        valid_loader = DataLoader(valid_dataset,
                                  batch_size=1,
                                  shuffle=False,
                                  num_workers=args.n_workers,
                                  pin_memory=True)

    return train_datasets, valid_datasets, train_loader, valid_loader


class Trainer(object):

    def __init__(self, args):
        self.train_datasets, \
        self.valid_datasets, \
        self.train_loader, \
        self.valid_loader = create_dataloaders(args)

        self.cfg = args

        # --------------Load and set model and optimizer-------------------------------------
        self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
        # self.device = torch.device('cpu')

        self.class_values = self.train_datasets[0].class_values
        self.ignore_label = VOID_VALUE if VOID_VALUE in self.class_values else 0
        print('Ignoring label value: %i' % self.ignore_label)

        self.classes = self.train_datasets[0].CLASSES
        self.non_bg_classes = np.asarray(self.classes)[np.asarray(self.class_values) != self.ignore_label]

        self.model = self.prepare_model()
        # Create adam optimizer
        self.optimizer = torch.optim.Adam(params=self.model.parameters(), lr=args.lr)

        # Dice/F1 score - https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient
        assert self.cfg.loss_fn in ['cross_entropy', 'dice', 'lovasz']
        if self.cfg.loss_fn == 'dice':
            self.criterion_fn = smp.losses.DiceLoss(mode='multiclass',
                                                    log_loss=True,
                                                    from_logits=True,
                                                    ignore_index=self.ignore_label)
        elif self.cfg.loss_fn == 'lovasz':
            self.criterion_fn = smp.losses.LovaszLoss(mode='multiclass',
                                                      from_logits=True,
                                                      ignore_index=self.ignore_label)
        elif self.cfg.loss_fn == 'cross_entropy':
            weights = torch.as_tensor([0.8, 0.2], device=self.device) if self.cfg.output == 'flexibility' else None
            self.criterion_fn = torch.nn.CrossEntropyLoss(ignore_index=self.ignore_label,
                                                          weight=weights)

        # IoU/Jaccard score - https://en.wikipedia.org/wiki/Jaccard_index
        # self.metric_fn = smp.utils.metrics.IoU(threshold=0.5)
        self.metric_fn = torchmetrics.JaccardIndex(num_classes=len(self.non_bg_classes),
                                                   ignore_index=self.ignore_label,
                                                   task='multiclass',
                                                   average='none').to(self.device)

        log_dir = '%s_lr_%g_bs_%d_%s_%s_labels_%s_%s' % \
                  (args.architecture, args.lr, args.batch_size,
                   '_'.join(args.datasets), '_'.join(self.train_datasets[0].fields),
                   self.train_datasets[0].output, time())
        self.tb_logger = SummaryWriter(log_dir=os.path.join(os.path.dirname(__file__), '../../config/tb_runs', log_dir))
        self.train_itr = 0
        self.val_itr = 0

    def __str__(self):
        return 'Training a model: %s\n' \
               'with batch size: %s\n' \
               'using input: %s\n' \
               'on datasets: %s\n' \
               'model output mode: %s\n' \
               'initial learning rate: %s' % \
               (self.cfg.architecture, self.cfg.batch_size, ' '.join(self.cfg.data_fields),
                ' '.join(self.cfg.datasets), self.cfg.output, self.cfg.lr)

    def prepare_model(self):
        if self.cfg.pretrained_weights is None:
            n_inputs = self.train_datasets[0][0][0].shape[0]
            n_classes = len(self.non_bg_classes)
            print('Model takes as input %i argument: %s' % (n_inputs, str(self.cfg.data_fields)))

            model = create_model(self.cfg.architecture, n_inputs, n_classes, pretrained_backbone=False)
        else:
            assert os.path.exists(self.cfg.pretrained_weights)
            model = torch.load(self.cfg.pretrained_weights)
        model = model.to(self.device)
        return model

    def compute_metric(self, pred, labels):
        # https://stackoverflow.com/questions/48260415/pytorch-how-to-compute-iou-jaccard-index-for-semantic-segmentation
        N, C, H, W = pred.shape
        assert labels.shape == (N, H, W)
        mask = labels != self.ignore_label

        pred = torch.softmax(pred, dim=1)
        pred = pred * mask.unsqueeze(1)
        labels = labels * mask

        ious = self.metric_fn(pred, labels)
        classes = self.non_bg_classes
        # assert len(classes) == len(ious)

        for i in range(len(ious)):
            # print('IOU for class %s: %f' % (self.train_datasets[0].CLASSES[i], ious[i]))
            self.tb_logger.add_scalar('IOU for class %s' % classes[i], ious[i], self.val_itr)
        iou = torch.mean(ious)

        return iou

    def train_epoch(self):
        losses_epoch = []
        for sample in tqdm(self.train_loader):
            inpt, labels = sample
            inpt, labels = inpt.to(self.device), labels.to(self.device)

            pred = self.model(inpt)['out']  # make prediction

            self.optimizer.zero_grad()
            loss = self.criterion_fn(pred, labels.long())  # Calculate loss
            loss.backward()  # Backpropagate loss
            self.optimizer.step()  # Apply gradient descent change to weight

            losses_epoch.append(loss.item())

            self.tb_logger.add_scalar('Train Loss (iter)', loss.item(), self.train_itr)
            self.train_itr += 1

        return np.mean(losses_epoch)

    def val_epoch(self):
        # validation epoch
        metrics_epoch = []
        losses_epoch = []
        for sample in tqdm(self.valid_loader):
            inpt, labels = sample
            inpt, labels = inpt.to(self.device), labels.to(self.device)

            with torch.no_grad():
                pred = self.model(inpt)['out']  # make prediction

                metric_sample = self.compute_metric(pred, labels)
                loss_val = self.criterion_fn(pred, labels.long())

            iou = metric_sample.cpu().numpy()
            metrics_epoch.append(iou)
            losses_epoch.append(loss_val.item())

            self.tb_logger.add_scalar('Val mIoU (iter)', iou, self.val_itr)
            self.tb_logger.add_scalar('Val Loss (iter)', loss_val.item(), self.val_itr)
            self.val_itr += 1

        metric_val = np.mean(metrics_epoch)
        loss_val = np.mean(losses_epoch)

        return loss_val, metric_val

    def test_model(self, dataset=None):
        self.model = self.model.eval()

        # Use the current trained model and visualize a prediction
        if dataset is None:
            ds_i = np.random.choice(range(len(self.valid_datasets)))
            dataset = self.valid_datasets[ds_i]

        inpt, label = dataset[np.random.choice(range(len(dataset)))]

        inpt = torch.from_numpy(inpt[None]).to(self.device)
        label = torch.from_numpy(label[None]).to(self.device)

        with torch.no_grad():
            pred = self.model(inpt)['out']

        pred = pred.squeeze(0).cpu().numpy()
        label = label.squeeze(0).cpu().numpy()

        color_pred = self.valid_datasets[0].label_to_color(pred)
        color_gt = self.valid_datasets[0].label_to_color(label)

        power = 16
        depth_img = np.copy(inpt.squeeze(0).cpu().numpy()[-1])  # depth
        depth_img[depth_img > 0] = depth_img[depth_img > 0] ** (1 / power)
        depth_img[depth_img > 0] = (depth_img[depth_img > 0] - depth_img[depth_img > 0].min()) / \
                                   (depth_img[depth_img > 0].max() - depth_img[depth_img > 0].min())

        self.tb_logger.add_image('Prediction', color_pred, dataformats='HWC')
        self.tb_logger.add_image('Ground truth', color_gt, dataformats='HWC')
        self.tb_logger.add_image('Depth image', depth_img, dataformats='HW')

        if self.cfg.vis_preds:
            color_pred_masked = color_pred.copy()
            color_pred_masked[label == 255] = color_gt[label == 255]

            label_flex = np.argmax(pred, axis=0) == dataset.mask_targets['flexible']
            depth_img_with_flex_points = (0.3 * depth_img + 0.7 * label_flex).astype("float")

            visualize_imgs(layout='columns',
                           prediction=color_pred,
                           prediction_masked=color_pred_masked,
                           ground_truth=color_gt,
                           # depth_img=depth_img,
                           depth_img=depth_img_with_flex_points,
                           )

            # visualize_cloud(xyz=dataset.scan.proj_xyz[::dataset.lidar_H_step].reshape((-1, 3)),
            #                 color=color_pred.reshape((-1, 3)))

    def train(self):
        print(self)

        max_metric = -np.Inf
        for epoch_n in tqdm(range(self.cfg.n_epochs)):
            print('Starting training epoch %i...' % epoch_n)
            # train epoch
            self.model = self.model.train()
            loss_train = self.train_epoch()

            print('Train loss at epoch %i: %f' % (epoch_n, float(loss_train)))
            self.tb_logger.add_scalar('Train Loss(epoch)', loss_train, epoch_n)

            print('Validation ...')
            self.model = self.model.eval()
            loss_val, metric_val = self.val_epoch()

            if not self.cfg.dont_save_models:
                # save better model
                if max_metric < metric_val:  # Save model weights
                    max_metric = metric_val

                if max_metric <= metric_val or epoch_n % 10 == 0:
                    h, w = self.train_datasets[0][0][1].shape[-2:]
                    name = '%s_lr_%g_bs_%d_epoch_%d_%s_%s_%sx%s_labels_%s_iou_%.3f.pth' % \
                           (self.cfg.architecture,
                            self.cfg.lr, self.cfg.batch_size, epoch_n,
                            '_'.join(self.cfg.datasets), '_'.join(self.train_datasets[0].fields),
                            h, w,
                            self.train_datasets[0].output, float(metric_val))

                    print("Saving Model:", name)
                    torch.save(self.model, os.path.join(os.path.dirname(__file__), name))

            print('Validation mIoU at epoch %i: %f' % (epoch_n, float(metric_val)))
            self.tb_logger.add_scalar('Val mIoU(epoch)', metric_val, epoch_n)
            self.tb_logger.add_scalar('Valid Loss(epoch)', loss_val, epoch_n)

            # change learning rate
            if epoch_n % 30 == 0 and epoch_n > 0:
                self.optimizer.param_groups[0]['lr'] /= 10.0
                print('Decrease decoder learning rate to %f !' % self.optimizer.param_groups[0]['lr'])

            self.test_model()

        self.tb_logger.close()


def main():
    args = parse_arguments()
    trainer = Trainer(args)
    trainer.train()


if __name__ == '__main__':
    main()