dlc2action.ssl.base_ssl

  1#
  2# Copyright 2020-2022 by A. Mathis Group and contributors. All rights reserved.
  3#
  4# This project and all its files are licensed under GNU AGPLv3 or later version. A copy is included in dlc2action/LICENSE.AGPL.
  5#
  6"""
  7Abstract class for defining SSL tasks
  8"""
  9
 10from typing import Callable, Tuple, Dict, Union
 11from abc import ABC, abstractmethod
 12import torch
 13from torch import nn
 14
 15
 16class SSLConstructor(ABC):
 17    """
 18    A base class for all SSL constructors
 19
 20    An SSL method is defined by three things: a *transformation* that maps a sample into SSL input and output,
 21    a neural net *module* that takes features as input and predicts SSL target, a *type* and a *loss function*.
 22    """
 23
 24    type = "none"
 25    """
 26    The `type` parameter defines interaction with the model:
 27    
 28    - `'ssl_input'`: a modification of the input data passes through the base network feature extraction module and the
 29    SSL module; it is returned as SSL output and compared to SSL target (or, if it is None, to the input data),
 30    - `'ssl_output'`:  the input data passes through the base network feature extraction module and the SSL module; it
 31    is returned as SSL output and compared to SSL target (or, if it is None, to the input data),
 32    - `'contrastive'`:  the input data and its modification pass through the base network feature extraction module and
 33    the SSL module; an (input results, modification results) tuple is returned as SSL output,
 34    - `'contrastive_2layers'`: the input data and its modification pass through the base network feature extraction module; 
 35    the output of the second feature extraction layer for the modified data goes through an SSL module and then, 
 36    optionally, that result and the first-level unmodified features pass another transformation; 
 37    an (input results, modified results) tuple is returned as SSL output,
 38    """
 39
 40    def __init__(self, *args, **kwargs) -> None:
 41        ...
 42
 43    @abstractmethod
 44    def transformation(self, sample_data: Dict) -> Tuple[Dict, Dict]:
 45        """
 46        Transform a sample feature dictionary into SSL input and target
 47
 48        Either input, target or both can be left as `None`. Transformers can be configured to replace `None` SSL targets
 49        with the input sample at runtime and/or to replace `None SSL` inputs with a new augmentation of the input sample.
 50        If the keys of the feature dictionaries are recognized by the transformer, they will be augmented before
 51        all features are stacked together.
 52
 53        Parameters
 54        ----------
 55        sample_data : dict
 56            a feature dictionary
 57
 58        Returns
 59        -------
 60        ssl_input : dict | torch.float('nan')
 61            a feature dictionary of SSL inputs
 62        ssl_target : dict | torch.float('nan')
 63            a feature dictionary of SSL targets
 64        """
 65
 66    @abstractmethod
 67    def loss(self, predicted: torch.Tensor, target: torch.Tensor) -> float:
 68        """
 69        Calculate the SSL loss
 70
 71        Parameters
 72        ----------
 73        predicted : torch.Tensor
 74            output of the SSL module
 75        target : torch.Tensor
 76            augmented and stacked SSL_target
 77
 78        Returns
 79        -------
 80        loss : float
 81            the loss value
 82        """
 83
 84    @abstractmethod
 85    def construct_module(self) -> nn.Module:
 86        """
 87        Construct the SSL module
 88
 89        Returns
 90        -------
 91        ssl_module : torch.nn.Module
 92            a neural net module that takes features extracted by a model's feature extractor as input and
 93            returns SSL output
 94        """
 95
 96
 97class EmptySSL(SSLConstructor):
 98    """
 99    Empty SSL class
100    """
101
102    def transformation(self, sample_data: Dict) -> Tuple:
103        """
104        Empty transformation
105        """
106
107        return (torch.tensor(float("nan")), torch.tensor(float("nan")))
108
109    def loss(self, predicted: torch.Tensor, target: torch.Tensor) -> float:
110        """
111        Empty loss
112        """
113
114        return 0
115
116    def construct_module(self) -> None:
117        """
118        Empty module
119        """
120
121        return None