train.py

import os
import sys
import copy
import argparse

import torch
from torch import optim
import torch.nn as nn

import mlflow.pytorch
from torch.utils.data import DataLoader
from torchvision.models import resnet18
import torchvision.transforms as T
from pytorch_lightning.metrics.functional import accuracy
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint

from utils.base import AuxLoss, WeightedLoss, display_mlflow_run_info, l2_regularization, str2bool, fetch_from_mlflow, get_name, data_loader_mean_and_std
from utils.dataset_utils import k_fold
from utils.augmentation import get_augmentation
from dataset import Subset, get_dataset

from processing.pipeline_numpy import RawProcessingPipeline
from processing.pipeline_torch import append_additive_layer, raw2rgb, RawToRGB, ParametrizedProcessing, NNProcessing

from model import log_tensor, resnet_model, LitModel, TrackImagesCallback

import segmentation_models_pytorch as smp

from utils.ssim import SSIM

# args to set up task
parser = argparse.ArgumentParser(description='classification_task')
parser.add_argument('--tracking_uri', type=str,
                    default='http://deplo-mlflo-1ssxo94f973sj-890390d809901dbf.elb.eu-central-1.amazonaws.com', help='URI of the mlflow server on AWS')
parser.add_argument('--processor_uri', type=str, default=None,
                    help='URI of the processing model (e.g. s3://mlflow-artifacts-821771080529/1/5fa754c566e3466690b1d309a476340f/artifacts/processing-model)')
parser.add_argument('--classifier_uri', type=str, default=None,
                    help='URI of the net (e.g. s3://mlflow-artifacts-821771080529/1/5fa754c566e3466690b1d309a476340f/artifacts/prediction-model)')
parser.add_argument('--state_dict_uri', type=str,
                    default=None, help='URI of the indices you want to load (e.g. s3://mlflow-artifacts-601883093460/7/4326da05aca54107be8c554de0674a14/artifacts/training')

parser.add_argument('--experiment_name', type=str,
                    default='classification learnable pipeline', help='Specify the experiment you are running, e.g. end2end segmentation')
parser.add_argument('--run_name', type=str,
                    default='test run', help='Specify the name of your run')

parser.add_argument('--log_model', type=str2bool, default=True, help='Enables model logging')
parser.add_argument('--save_locally', action='store_true',
                    help='Model will be saved locally if action is taken')   # TODO: bypass mlflow

parser.add_argument('--track_processing', action='store_true',
                    help='Save images after each trasformation of the pipeline for the test set')
parser.add_argument('--track_processing_gradients', action='store_true',
                    help='Save images of gradients after each trasformation of the pipeline for the test set')
parser.add_argument('--track_save_tensors', action='store_true',
                    help='Save the torch tensors after each trasformation of the pipeline for the test set')
parser.add_argument('--track_predictions', action='store_true',
                    help='Save images after each trasformation of the pipeline for the test set + input gradient')
parser.add_argument('--track_n_images', default=5,
                    help='Track the n first elements of dataset. Only used for args.track_processing=True')
parser.add_argument('--track_every_epoch', action='store_true', help='Track images every epoch or once after training')

# args to create dataset
parser.add_argument('--seed', type=int, default=1, help='Global seed')
parser.add_argument('--dataset', type=str, default='Microscopy',
                    choices=['Drone', 'DroneSegmentation', 'Microscopy'], help='Select dataset')

parser.add_argument('--n_splits', type=int, default=1, help='Number of splits used for training')
parser.add_argument('--train_size', type=float, default=0.8, help='Fraction of training points in dataset')

# args for training
parser.add_argument('--lr', type=float, default=1e-5, help='learning rate used for training')
parser.add_argument('--epochs', type=int, default=3, help='numper of epochs')
parser.add_argument('--batch_size', type=int, default=32, help='Training batch size')
parser.add_argument('--augmentation', type=str, default='none',
                    choices=['none', 'weak', 'strong'], help='Applies augmentation to training')
parser.add_argument('--check_val_every_n_epoch', type=int, default=1)

# args to specify the processing
parser.add_argument('--processing_mode', type=str, default='parametrized',
                    choices=['parametrized', 'static', 'neural_network', 'none'],
                    help='Which type of raw to rgb processing should be used')

# args to specify model
parser.add_argument('--classifier_network', type=str, default='ResNet18', choices=['ResNet18', 'ResNet34', 'Resnet50'],
                    help='Type of pretrained network')
parser.add_argument('--classifier_pretrained', action='store_true',
                    help='Whether to use a pre-trained model or not')
parser.add_argument('--smp_encoder', type=str, default='resnet34', help='segmentation models pytorch encoder')

parser.add_argument('--freeze_processor', action='store_true', help='Freeze raw to rgb processing model weights')
parser.add_argument('--freeze_classifier', action='store_true', help='Freeze classification model weights')

# args to specify static pipeline transformations
parser.add_argument('--sp_debayer', type=str, default='bilinear',
                    choices=['bilinear', 'malvar2004', 'menon2007'], help='Specify algorithm used as debayer')
parser.add_argument('--sp_sharpening', type=str, default='sharpening_filter',
                    choices=['sharpening_filter', 'unsharp_masking'], help='Specify algorithm used for sharpening')
parser.add_argument('--sp_denoising', type=str, default='gaussian_denoising',
                    choices=['gaussian_denoising', 'median_denoising', 'fft_denoising'], help='Specify algorithm used for denoising')

# args to choose training mode
parser.add_argument('--adv_training', action='store_true', help='Enable adversarial training')
parser.add_argument('--adv_aux_weight', type=float, default=1, help='Weighting of the adversarial auxilliary loss')
parser.add_argument('--adv_aux_loss', type=str, default='ssim', choices=['l2', 'ssim'],
                    help='Type of adversarial auxilliary regularization loss')
parser.add_argument('--adv_noise_layer', action='store_true', help='Adds an additive layer to Parametrized Processing')
parser.add_argument('--adv_track_differences', action='store_true', help='Save difference to default pipeline')
parser.add_argument('--adv_parameters', choices=['all', 'black_level', 'white_balance',
                                                 'colour_correction', 'gamma_correct', 'sharpening_filter', 'gaussian_blur', 'additive_layer'],
                    help='Target individual parameters for adversarial training.')

parser.add_argument('--cache_downloaded_models', type=str2bool, default=True)

parser.add_argument('--test_run', action='store_true')


args = parser.parse_args()

os.makedirs('results', exist_ok=True)


def run_train(args):

    print(args)

    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
    training_mode = 'adversarial' if args.adv_training else 'default'

    # set tracking uri, this is the address of the mlflow server where light experimental data will be stored
    mlflow.set_tracking_uri(args.tracking_uri)
    mlflow.set_experiment(args.experiment_name)
    os.environ['AWS_ACCESS_KEY_ID'] = '#TODO: fill in your aws access key id for mlflow server here'
    os.environ['AWS_SECRET_ACCESS_KEY'] = '#TODO: fill in your aws secret access key for mlflow server here'

    dataset = get_dataset(args.dataset)

    print(f'dataset: {type(dataset).__name__}[{len(dataset)}]')
    print(f'task: {dataset.task}')
    print(f'mode: {training_mode} training')
    print(f'# cross-validation subsets: {args.n_splits}')
    pl.seed_everything(args.seed)
    idxs_kfold = k_fold(dataset, n_splits=args.n_splits, seed=args.seed, train_size=args.train_size)

    # start mlflow parent run for k-fold validation (optional)
    with mlflow.start_run(run_name=args.run_name) as parent_run:

        # start mlflow child run
        for k_iter, (train_indices, valid_indices) in enumerate(idxs_kfold):

            print(f'K_fold subset: {k_iter+1}/{args.n_splits}')

            if args.processing_mode == 'static':
                # only needed if processor outputs should be normalized (might help for classifier training / testing against torch pipeline)
                if args.dataset == 'Drone' or args.dataset == 'DroneSegmentation':
                    mean = torch.tensor([0.35, 0.36, 0.35])
                    std = torch.tensor([0.12, 0.11, 0.12])
                elif args.dataset == 'Microscopy':
                    mean = torch.tensor([0.91, 0.84, 0.94])
                    std = torch.tensor([0.08, 0.12, 0.05])

                # numpy pipeline doesn't use torch batched transformations. Transformations are applied individually to dataloader
                dataset.transform = T.Compose([RawProcessingPipeline(
                    camera_parameters=dataset.camera_parameters,
                    debayer=args.sp_debayer,
                    sharpening=args.sp_sharpening,
                    denoising=args.sp_denoising,
                ),
                    T.Normalize(mean, std)
                ])

                processor = nn.Identity()

            # fetch processor from mlflow
            if args.processor_uri is not None and args.processing_mode != 'none':
                print('Fetching processor: ', end='')
                processor = fetch_from_mlflow(args.processor_uri, type='processor',
                                              use_cache=args.cache_downloaded_models)
            else:
                print(f'processing_mode: {args.processing_mode}')
                normalize_mosaic = None   # normalize after raw has been transformed to rgb image via raw2rgb
                # not strictly necessary, but for processing_mode=='none' this will ensure normalized outputs for the classifier
                # and for processing_mode=='neural_network', the processing segmentation model receives normalized inputs
                # could be evaded via an additional batchnorm!
                # XXX
                if args.dataset == 'Microscopy':
                    mosaic_mean = [0.5663, 0.1401, 0.0731]
                    mosaic_std = [0.097, 0.0423, 0.008]
                    normalize_mosaic = T.Normalize(mosaic_mean, mosaic_std)

                # track individual processing steps for visualization
                track_stages = args.track_processing or args.track_processing_gradients
                if args.processing_mode == 'parametrized':
                    processor = ParametrizedProcessing(
                        camera_parameters=dataset.camera_parameters, track_stages=track_stages, batch_norm_output=True)

                elif args.processing_mode == 'neural_network':
                    processor = NNProcessing(track_stages=track_stages,
                                             normalize_mosaic=normalize_mosaic, batch_norm_output=True)
                elif args.processing_mode == 'none':
                    processor = RawToRGB(reduce_size=True, out_channels=3, track_stages=track_stages,
                                         normalize_mosaic=normalize_mosaic)

            if args.classifier_uri:  # fetch classifier from mlflow
                print('Fetching classifier: ', end='')
                classifier = fetch_from_mlflow(args.classifier_uri, type='classifier',
                                               use_cache=args.cache_downloaded_models)
            else:
                if dataset.task == 'classification':
                    classifier = resnet_model(
                        model=args.classifier_network,
                        pretrained=args.classifier_pretrained,
                        in_channels=3,
                        fc_out_features=len(dataset.classes)
                    )
                else:
                    classifier = smp.UnetPlusPlus(
                        encoder_name=args.smp_encoder,
                        encoder_depth=5,
                        encoder_weights='imagenet',
                        in_channels=3,
                        classes=1,
                        activation=None,
                    )

            if args.freeze_processor and len(list(iter(processor.parameters()))) == 0:
                print('Note: freezing processor without parameters.')
            assert not (args.freeze_processor and args.freeze_classifier), 'Likely no parameters to train.'

            if dataset.task == 'classification':
                loss = nn.CrossEntropyLoss()
                metrics = [accuracy]
            else:
                # loss = utils.base.smp_get_loss(args.smp_loss)    # XXX: add other losses to args.smp_loss
                loss = smp.losses.DiceLoss(mode='binary', from_logits=True)
                metrics = [smp.utils.metrics.IoU()]

            loss_aux = None

            if args.adv_training:   # setup for failure mode search

                assert args.processing_mode == 'parametrized', f"Processing mode ({args.processing_mode}) should be set to 'parametrized' for adversarial training"
                assert args.freeze_classifier, 'Classifier should be frozen for adversarial training'
                assert not args.freeze_processor, 'Processor should not be frozen for adversarial training'

                # copy, so that regularization in rgb space between adversarial and original processor can be computed
                processor_default = copy.deepcopy(processor)
                processor_default.track_stages = args.track_processing
                processor_default.eval()
                processor_default.to(DEVICE)

                for p in processor_default.parameters():
                    p.requires_grad = False

                if args.adv_noise_layer:    # optional additional "noise" layer in processor
                    append_additive_layer(processor)

                if args.adv_aux_loss == 'l2':
                    regularization = l2_regularization
                elif args.adv_aux_loss == 'ssim':
                    regularization = SSIM(window_size=11)
                else:
                    NotImplementedError(args.adv_aux_loss)

                loss = WeightedLoss(loss=loss, weight=-1)

                loss_aux = AuxLoss(
                    loss_aux=regularization,
                    processor_adv=processor,
                    processor_default=processor_default,
                    weight=args.adv_aux_weight,
                )

            augmentation = get_augmentation(args.augmentation)

            model = LitModel(
                classifier=classifier,
                processor=processor,
                loss=loss,
                lr=args.lr,
                loss_aux=loss_aux,
                adv_training=args.adv_training,
                adv_parameters=args.adv_parameters,
                metrics=metrics,
                augmentation=augmentation,
                is_segmentation_task=dataset.task == 'segmentation',
                freeze_classifier=args.freeze_classifier,
                freeze_processor=args.freeze_processor,
            )

            state_dict = vars(args).copy()

            # get train_set_dict
            if args.state_dict_uri:
                state_dict = mlflow.pytorch.load_state_dict(args.state_dict_uri)
                train_indices = state_dict['train_indices']
                valid_indices = state_dict['valid_indices']

            track_indices = list(range(args.track_n_images))

            if dataset.task == 'classification':
                state_dict['classes'] = dataset.classes
            state_dict['device'] = DEVICE
            state_dict['train_indices'] = train_indices
            state_dict['valid_indices'] = valid_indices
            state_dict['elements in train set'] = len(train_indices)
            state_dict['elements in test set'] = len(valid_indices)

            if args.test_run:
                train_indices = train_indices[:args.batch_size]
                valid_indices = valid_indices[:args.batch_size]

            train_set = Subset(dataset, indices=train_indices)
            valid_set = Subset(dataset, indices=valid_indices)
            track_set = Subset(dataset, indices=track_indices)

            train_loader = DataLoader(train_set, batch_size=args.batch_size, num_workers=16, shuffle=True)
            valid_loader = DataLoader(valid_set, batch_size=args.batch_size, num_workers=16, shuffle=False)
            track_loader = DataLoader(track_set, batch_size=args.batch_size, num_workers=16, shuffle=False)

            with mlflow.start_run(run_name=f"{args.run_name}_{k_iter}", nested=True) as child_run:

                if k_iter == 0:
                    display_mlflow_run_info(child_run)

                mlflow.pytorch.log_state_dict(state_dict, artifact_path=None)

                hparams = {
                    'dataset': args.dataset,
                    'processing_mode': args.processing_mode,
                    'training_mode': training_mode,
                }
                if training_mode == 'adversarial':
                    hparams['adv_aux_weight'] = args.adv_aux_weight
                    hparams['adv_aux_loss'] = args.adv_aux_loss

                mlflow.log_params(hparams)

                with open('results/state_dict.txt', 'w') as f:
                    f.write('python ' + ' '.join(sys.argv) + '\n')
                    f.write('\n'.join([f'{k}={v}' for k, v in state_dict.items()]))
                mlflow.log_artifact('results/state_dict.txt', artifact_path=None)

                mlf_logger = pl.loggers.MLFlowLogger(experiment_name=args.experiment_name,
                                                     tracking_uri=args.tracking_uri,)
                mlf_logger._run_id = child_run.info.run_id

                reference_processor = processor_default if args.adv_training and args.adv_track_differences else None

                callbacks = []
                if args.track_processing:
                    callbacks += [TrackImagesCallback(track_loader,
                                                      reference_processor,
                                                      track_every_epoch=args.track_every_epoch,
                                                      track_processing=args.track_processing,
                                                      track_gradients=args.track_processing_gradients,
                                                      track_predictions=args.track_predictions,
                                                      save_tensors=args.track_save_tensors)]

                trainer = pl.Trainer(
                    gpus=1 if DEVICE == 'cuda' else 0,
                    min_epochs=args.epochs,
                    max_epochs=args.epochs,
                    logger=mlf_logger,
                    callbacks=callbacks,
                    check_val_every_n_epoch=args.check_val_every_n_epoch,
                )

                if args.log_model:
                    mlflow.pytorch.autolog(log_every_n_epoch=10)
                    print(f'model_uri="{mlflow.get_artifact_uri()}/model"')

                t = trainer.fit(
                    model,
                    train_dataloader=train_loader,
                    val_dataloaders=valid_loader,
                )

    globals().update(locals())  # for convenient access

    return model


if __name__ == '__main__':
    model = run_train(args)