cleanup

model.py

@@ -3,7 +3,7 @@ from collections import defaultdict
 
 import torch
 import torch.optim
-from torchvision.models import resnet18
+from torchvision.models import resnet18, resnet34, resnet50
 from torchvision.utils import make_grid, save_image
 import torch.nn.functional as F
 
@@ -12,10 +12,13 @@ import pytorch_lightning as pl
 import mlflow.pytorch
 
 
-def resnet_model(model=resnet18, pretrained=True, in_channels=3, fc_out_features=2):
-    resnet = model(pretrained=pretrained)
-    # if not pretrained:  # TODO: add case for in_channels=4
-    #     resnet.conv1 = torch.nn.Conv2d(channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
+def resnet_model(model='resnet18', pretrained=True, in_channels=3, fc_out_features=2):
+    if model.lower() == 'resnet18':
+        resnet = resnet18(pretrained=pretrained)
+    if model.lower() == 'resnet34':
+        resnet = resnet34(pretrained=pretrained)
+    if model.lower() == 'resnet50':
+        resnet = resnet50(pretrained=pretrained)
     resnet.fc = torch.nn.Linear(in_features=512, out_features=fc_out_features, bias=True)
     return resnet
 
@@ -81,10 +84,6 @@ class LitModel(pl.LightningModule):
 
     def update_step(self, batch, step_name):
         x, y = batch
-        # debug(self.processor)
-        # debug(self.processor.parameters())
-        # debug.pause()
-        # print('type', type(self.processor).__name__)
 
         logits = self(x)
 

processing/pipeline_torch.py

@@ -43,6 +43,18 @@ DEFAULT_CAMERA_PARAMS = (
 
 
 class RawToRGB(nn.Module):
+    """transforms a raw image with 1 channel to rgb with 3 channels 
+
+    Args:
+        reduce_size (bool, optional): if False, the output image will have the same height and width 
+            as the raw input, i.e. (B, C, H, W), empty values are filled with zeros.
+            if True, the output dimensions are reduced by half (B, C, H//2, W//2), 
+        out_channels (int, optional): number of output channels. One of {3, 4}.
+            for 3 channels, the two green channels are averaged.
+        track_stages (bool, optional): whether or not to retain intermediary steps in processing
+        normalize_mosaic (function, optional): applies normalization transformation to rgb image
+    """
+
     def __init__(self, reduce_size=True, out_channels=3, track_stages=False, normalize_mosaic=None):
         super().__init__()
         self.stages = None
@@ -71,6 +83,14 @@ class RawToRGB(nn.Module):
 
 
 class NNProcessing(nn.Module):
+    """Transforms raw images to processed rgb via a segmentation Unet
+
+    Args:
+        track_stages (bool, optional): whether or not to retain intermediary steps in processing
+        normalize_mosaic (function, optional): applies normalization transformation to rgb image
+        batch_norm_output (bool, optional): adds a BatchNorm layer to the end of the processing
+    """
+
     def __init__(self, track_stages=False, normalize_mosaic=None, batch_norm_output=True):
         super().__init__()
         self.stages = None
@@ -89,7 +109,7 @@ class NNProcessing(nn.Module):
     def forward(self, raw):
         self.stages = {}
         self.buffer = {}
-        # self.stages['raw'] = raw
+
         rgb = raw2rgb(raw)
         if self.normalize_mosaic:
             rgb = self.normalize_mosaic(rgb)
@@ -108,18 +128,29 @@ class NNProcessing(nn.Module):
         return rgb
 
 
-def add_additive_layer(processor):
+def append_additive_layer(processor):
     processor.additive_layer = nn.Parameter(torch.zeros((1, 3, 256, 256)))
     # processor.additive_layer = nn.Parameter(0.001 * torch.randn((1, 3, 256, 256)))
 
 
 class ParametrizedProcessing(nn.Module):
-    def __init__(self, camera_parameters, track_stages=False, batch_norm_output=True):
+    """Differentiable processing pipeline via torch transformations
+
+    Args:
+        camera_parameters (tuple(list), optional): applies given camera parameters in processing
+        track_stages (bool, optional): whether or not to retain intermediary steps in processing
+        batch_norm_output (bool, optional): adds a BatchNorm layer to the end of the processing
+    """
+
+    def __init__(self, camera_parameters=None, track_stages=False, batch_norm_output=True):
         super().__init__()
         self.stages = None
         self.buffer = None
         self.track_stages = track_stages
 
+        if camera_parameters is None:
+            camera_parameters = DEFAULT_CAMERA_PARAMS
+
         black_level, white_balance, colour_matrix = camera_parameters
 
         self.black_level = nn.Parameter(torch.as_tensor(black_level))
@@ -197,8 +228,11 @@ class ParametrizedProcessing(nn.Module):
 
 
 class Debayer(nn.Conv2d):
+    """Separates the mosaiced raw image into its channels and interpolates bilinearly. Output is of same size as input.
+    """
+
     def __init__(self):
-        super().__init__(3, 3, kernel_size=3, padding=1, padding_mode='reflect', bias=False)    # default_pipeline uses 'replicate'
+        super().__init__(3, 3, kernel_size=3, padding=1, padding_mode='reflect', bias=False)    # pipeline_numpy uses 'replicate'
         self.weight.data.fill_(0)
         self.weight.data[0, 0] = K_RB.clone()
         self.weight.data[1, 1] = K_G.clone()
@@ -206,15 +240,16 @@ class Debayer(nn.Conv2d):
 
 
 def raw2rgb(raw, black_level=None, reduce_size=True, out_channels=3):
-    """transform raw image with 1 channel to rgb with 3 channels 
+    """Transforms a raw image with 1 channel to rgb with 3 channels 
+
     Args:
         raw (Tensor): raw Tensor of shape (B, H, W)
         black_level (iterable, optional): RGGB black level values to subtract
         reduce_size (bool, optional): if False, the output image will have the same height and width 
             as the raw input, i.e. (B, C, H, W), empty values are filled with zeros.
             if True, the output dimensions are reduced by half (B, C, H//2, W//2), 
-            the two green channels are averaged.
         out_channels (int, optional): number of output channels. One of {3, 4}.
+            The two green channels are averaged if out_channels == 3.
     """
     assert out_channels in [3, 4]
     if black_level is None:

train.py

@@ -15,14 +15,14 @@ from pytorch_lightning.metrics.functional import accuracy
 import pytorch_lightning as pl
 from pytorch_lightning.callbacks import ModelCheckpoint
 
-from utils.base import display_mlflow_run_info, str2bool, fetch_from_mlflow, get_name, data_loader_mean_and_std
+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.debug import debug
 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 add_additive_layer, raw2rgb, RawToRGB, ParametrizedProcessing, NNProcessing
+from processing.pipeline_torch import append_additive_layer, raw2rgb, RawToRGB, ParametrizedProcessing, NNProcessing
 
 from model import log_tensor, resnet_model, LitModel, TrackImagesCallback
 
@@ -31,89 +31,88 @@ 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,
+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,
+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,
+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,
+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,
+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',
+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',
+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',
+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',
+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',
+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,
+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')
+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('--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')
+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("--augmentation_on_valid_epoch", action='store_true',
-                    help='Track images every epoch or once after training')  # TODO: implement, actually should be disabled by default for 'val' and 'test
-parser.add_argument("--check_val_every_n_epoch", type=int, default=1)
+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")
+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',
-                    help='Type of pretrained network')  # TODO: implement different choices
-parser.add_argument("--classifier_pretrained", action='store_true',
+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 model encoder')
+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")
+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")
+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_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'])
+                                                 '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('--cache_downloaded_models', type=str2bool, default=True)
 
 parser.add_argument('--test_run', action='store_true')
 
@@ -133,10 +132,8 @@ def run_train(args):
     # 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
+    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)
 
@@ -147,52 +144,57 @@ def run_train(args):
     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:
 
-        for k_iter, idxs in enumerate(idxs_kfold):
+        # 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}")
+            print(f'K_fold subset: {k_iter+1}/{args.n_splits}')
 
             if args.processing_mode == 'static':
-                if args.dataset == "Drone" or args.dataset == "DroneSegmentation":
+                # 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":
+                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)])
-                # XXX: Not clean
+                ),
+                    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='')
-                model = fetch_from_mlflow(args.processor_uri, use_cache=args.cache_downloaded_models)
-                processor = model.processor
-                for param in processor.parameters():
-                    param.requires_grad = True
-                model.processor = None
-                del model
+                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 passed to raw2rgb
-                if args.dataset == "Microscopy":
+                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,
-                        # noise_layer=args.adv_noise_layer, # this has to be added manually afterwards for when a model is loaded that doesn't have one yet
-                    )
+                        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,
@@ -201,22 +203,19 @@ def run_train(args):
                     processor = RawToRGB(reduce_size=True, out_channels=3, track_stages=track_stages,
                                          normalize_mosaic=normalize_mosaic)
 
-            if args.classifier_uri:  # fetch classifier
+            if args.classifier_uri:  # fetch classifier from mlflow
                 print('Fetching classifier: ', end='')
-                model = fetch_from_mlflow(args.classifier_uri, use_cache=args.cache_downloaded_models)
-                classifier = model.classifier
-                model.classifier = None
-                del model
+                classifier = fetch_from_mlflow(args.classifier_uri, type='classifier',
+                                               use_cache=args.cache_downloaded_models)
             else:
                 if dataset.task == 'classification':
                     classifier = resnet_model(
-                        model=resnet18,
+                        model=args.classifier_network,
                         pretrained=args.classifier_pretrained,
                         in_channels=3,
                         fc_out_features=len(dataset.classes)
                     )
                 else:
-                    # XXX: add other network choices to args.smp_network (FPN) and args.network
                     classifier = smp.UnetPlusPlus(
                         encoder_name=args.smp_encoder,
                         encoder_depth=5,
@@ -240,26 +239,23 @@ def run_train(args):
 
             loss_aux = None
 
-            if args.adv_training:
+            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"
+                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)
-                # debug(processor_default)
+
                 for p in processor_default.parameters():
                     p.requires_grad = False
 
-                if args.adv_noise_layer:
-                    add_additive_layer(processor)
-
-                def l2_regularization(x, y):
-                    return ((x - y) ** 2).sum()
-                    # return (x - y).norm()
+                if args.adv_noise_layer:    # optional additional "noise" layer in processor
+                    append_additive_layer(processor)
 
                 if args.adv_aux_loss == 'l2':
                     regularization = l2_regularization
@@ -268,34 +264,12 @@ def run_train(args):
                 else:
                     NotImplementedError(args.adv_aux_loss)
 
-                class AuxLoss(nn.Module):
-                    def __init__(self, loss_aux, weight=1):
-                        super().__init__()
-                        self.loss_aux = loss_aux
-                        self.weight = weight
-
-                    def forward(self, x):
-                        with torch.no_grad():
-                            x_reference = processor_default(x)
-                        x_processed = processor.buffer['processed_rgb']
-                        return self.weight * self.loss_aux(x_reference, x_processed)
-
-                class WeightedLoss(nn.Module):
-                    def __init__(self, loss, weight=1):
-                        super().__init__()
-                        self.loss = loss
-                        self.weight = weight
-
-                    def forward(self, x, y):
-                        return self.weight * self.loss(x, y)
-
-                    def __repr__(self):
-                        return f'{self.weight} * {get_name(self.loss)}'
-
                 loss = WeightedLoss(loss=loss, weight=-1)
-                # loss = WeightedLoss(loss=nn.CrossEntropyLoss(), weight=0)
+
                 loss_aux = AuxLoss(
                     loss_aux=regularization,
+                    processor_adv=processor,
+                    processor_default=processor_default,
                     weight=args.adv_aux_weight,
                 )
 
@@ -316,15 +290,13 @@ def run_train(args):
                 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']
-            else:
-                train_indices = idxs[0]
-                valid_indices = idxs[1]
-                state_dict = vars(args).copy()
 
             track_indices = list(range(args.track_n_images))
 
@@ -350,8 +322,6 @@ def run_train(args):
 
             with mlflow.start_run(run_name=f"{args.run_name}_{k_iter}", nested=True) as child_run:
 
-                # mlflow.pytorch.autolog(silent=True)
-
                 if k_iter == 0:
                     display_mlflow_run_info(child_run)
 
@@ -389,14 +359,6 @@ def run_train(args):
                                                       track_predictions=args.track_predictions,
                                                       save_tensors=args.track_save_tensors)]
 
-                # if True: #args.save_best:
-                #    if dataset.task == 'classification':
-                    #checkpoint_callback = ModelCheckpoint(pathmonitor="val_accuracy", mode='max')
-                #        checkpoint_callback = ModelCheckpoint(dirpath=args.tracking_uri, save_top_k=1, verbose=True, monitor="val_accuracy", mode="max") #dirpath=args.tracking_uri,
-                #    else:
-                #        checkpoint_callback = ModelCheckpoint(monitor="val_iou_score")
-                #callbacks += [checkpoint_callback]
-
                 trainer = pl.Trainer(
                     gpus=1 if DEVICE == 'cuda' else 0,
                     min_epochs=args.epochs,
@@ -404,7 +366,6 @@ def run_train(args):
                     logger=mlf_logger,
                     callbacks=callbacks,
                     check_val_every_n_epoch=args.check_val_every_n_epoch,
-                    # checkpoint_callback=True,
                 )
 
                 if args.log_model:

utils/base.py

@@ -18,14 +18,7 @@ from b2sdk.v1 import *
 
 import argparse
 
-
-class SmartFormatter(argparse.HelpFormatter):
-      
-    def _split_lines(self, text, width):
-        if text.startswith('R|'):
-            return text[2:].splitlines()  
-        # this is the RawTextHelpFormatter._split_lines
-        return argparse.HelpFormatter._split_lines(self, text, width)
+from torch import nn
 
 
 def str2bool(string):
@@ -193,6 +186,7 @@ def b2_download_folder(b2_dir, local_dir, force_download=False, mirror_folder=Tr
 def get_name(obj):
     return obj.__name__ if hasattr(obj, '__name__') else type(obj).__name__
 
+
 def get_mlflow_model_by_name(experiment_name, run_name,
                              tracking_uri="http://deplo-mlflo-1ssxo94f973sj-890390d809901dbf.elb.eu-central-1.amazonaws.com",
                              download_model=True):
@@ -234,6 +228,7 @@ def get_mlflow_model_by_name(experiment_name, run_name,
 
     return state_dict, model
 
+
 def data_loader_mean_and_std(data_loader, transform=None):
     means = []
     stds = []
@@ -244,23 +239,35 @@ def data_loader_mean_and_std(data_loader, transform=None):
         stds.append(x.std(dim=(0, 2, 3)).unsqueeze(0))
     return torch.cat(means).mean(dim=0), torch.cat(stds).mean(dim=0)
 
+
 def fetch_runs_list_mlflow(experiment):
     runs = mlflow.search_runs(experiment.experiment_id)
     runs.to_pickle('cache/runs_names.pkl')  # where to save it, usually as a .pkl
     return runs
 
-def fetch_from_mlflow(uri, use_cache=True, download_model=True):
+
+def fetch_from_mlflow(uri, type='', use_cache=True, download_model=True):
     cache_loc = os.path.join('cache', uri.split('//')[1]) + '.pt'
     if use_cache and os.path.exists(cache_loc):
         print(f'loading cached model from {cache_loc} ...')
-        return torch.load(cache_loc)
+        model = torch.load(cache_loc)
     else:
         print(f'fetching model from {uri} ...')
         model = mlflow.pytorch.load_model(uri)
         os.makedirs(os.path.dirname(cache_loc), exist_ok=True)
         if download_model:
             torch.save(model, cache_loc, pickle_module=mlflow.pytorch.pickle_module)
-        return model
+    if type == 'processor':
+        processor = model.processor
+        model.processor = None
+        del model   # free up memory space
+        return processor
+    if type == 'classifier':
+        classifier = model.classifier
+        model.classifier = None
+        del model   # free up memory space
+        return classifier
+    return model
 
 
 def display_mlflow_run_info(run):
@@ -315,16 +322,50 @@ def get_train_test_indices_drone(df, frac, seed=None):
     return train_indices, test_indices
 
 
-def smp_get_loss(loss):
-    if loss == "Dice":
-        return smp.losses.DiceLoss(mode='binary', from_logits=True)
-    if loss == "BCE":
-        return nn.BCELoss()
-    elif loss == "BCEWithLogits":
-        return smp.losses.BCEWithLogitsLoss()
-    elif loss == "DicyBCE":
-        from pytorch_toolbelt import losses as ptbl
-        return ptbl.JointLoss(ptbl.DiceLoss(mode='binary', from_logits=False),
-                              nn.BCELoss(),
-                              first_weight=args.dice_weight,
-                              second_weight=args.bce_weight)
+# def smp_get_loss(loss):
+#     if loss == "Dice":
+#         return smp.losses.DiceLoss(mode='binary', from_logits=True)
+#     if loss == "BCE":
+#         return nn.BCELoss()
+#     elif loss == "BCEWithLogits":
+#         return smp.losses.BCEWithLogitsLoss()
+#     elif loss == "DicyBCE":
+#         from pytorch_toolbelt import losses as ptbl
+#         return ptbl.JointLoss(ptbl.DiceLoss(mode='binary', from_logits=False),
+#                               nn.BCELoss(),
+#                               first_weight=args.dice_weight,
+#                               second_weight=args.bce_weight)
+
+
+# Adversarial setup
+
+def l2_regularization(x, y):
+    return ((x - y) ** 2).sum()
+
+
+class AuxLoss(nn.Module):
+    def __init__(self, loss_aux, processor_adv, processor_default, weight=1):
+        super().__init__()
+        self.loss_aux = loss_aux
+        self.weight = weight
+        self.processor_adv = processor_adv
+        self.processor_default = processor_default
+
+    def forward(self, x):
+        with torch.no_grad():
+            x_reference = self.processor_default(x)
+        x_processed = self.processor.buffer['processed_rgb']
+        return self.weight * self.loss_aux(x_reference, x_processed)
+
+
+class WeightedLoss(nn.Module):
+    def __init__(self, loss, weight=1):
+        super().__init__()
+        self.loss = loss
+        self.weight = weight
+
+    def forward(self, x, y):
+        return self.weight * self.loss(x, y)
+
+    def __repr__(self):
+        return f'{self.weight} * {get_name(self.loss)}'