neutral_ad.py

# Neural Transformation Learning for Anomaly Detection (NeuTraLAD) - a self-supervised method for anomaly detection
# Copyright (c) 2022 Robert Bosch GmbH
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

from typing import Dict, List, Optional, Sequence, Tuple

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from timesead.models.common import AnomalyDetector
from timesead.models import BaseModel
from timesead.optim.loss import Loss
from timesead.utils.utils import pack_tuple
from timesead_ext.models.transforms import (
    Transform,
    TransformBank,
    freq_cfg as default_freq_cfg,
    group_cfg as default_group_cfg,
    invertible_cfg as default_invertible_cfg,
    make_freq_family,
    make_group_family,
    make_invertible_family,
)


class ResTrans1DBlock(torch.nn.Module):
    def __init__(self, channel: int, bias: bool = False):
        super().__init__()
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv1d(channel, channel, 3, 1, 1, bias=bias)
        self.in1 = nn.InstanceNorm1d(channel, affine=bias)
        self.conv2 = nn.Conv1d(channel, channel, 3, 1, 1, bias=bias)
        self.in2 = nn.InstanceNorm1d(channel, affine=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        out = self.relu(self.in1(self.conv1(x)))
        out = self.in2(self.conv2(out))
        out = out + residual
        out = self.relu(out)
        return out


class ConvLayer(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int,
                 dilation: int = 1, bias: bool = False):
        super().__init__()
        padding = dilation * (kernel_size // 2)
        self.reflection_pad = nn.ReflectionPad1d(padding)
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, stride, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.reflection_pad(x)
        out = self.conv1d(out)
        return out


class SeqTransformNet(Transform):
    def __init__(self, x_dim: int, hdim: int, num_layers: int):
        super().__init__()
        self.relu = nn.ReLU()
        self.conv1 = ConvLayer(x_dim, hdim, 3, 1, bias=False)
        self.in1 = nn.InstanceNorm1d(hdim, affine=False)
        res_blocks = []
        for _ in range(num_layers - 2):
            res_blocks.append(ResTrans1DBlock(hdim, False))
        self.res = nn.Sequential(*res_blocks)
        self.conv2 = ConvLayer(hdim, x_dim, 3, 1, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.relu(self.in1(self.conv1(x)))
        out = self.res(out)
        out = self.conv2(out)
        return out


class ResBlock(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, conv_param=None, downsample=None,
                 batchnorm: bool = False, bias: bool = False):
        super().__init__()

        self.conv1 = nn.Conv1d(in_dim, in_dim, 1, 1, 0, bias=bias)
        if conv_param is not None:
            self.conv2 = nn.Conv1d(in_dim, in_dim, conv_param[0], conv_param[1], conv_param[2], bias=bias)
        else:
            self.conv2 = nn.Conv1d(in_dim, in_dim, 3, 1, 1, bias=bias)

        self.conv3 = nn.Conv1d(in_dim, out_dim, 1, 1, 0, bias=bias)
        if batchnorm:
            self.bn1 = nn.BatchNorm1d(in_dim)
            self.bn2 = nn.BatchNorm1d(in_dim)
            self.bn3 = nn.BatchNorm1d(out_dim)
            if downsample:
                self.bn4 = nn.BatchNorm1d(out_dim)

        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.batchnorm = batchnorm

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x

        out = self.conv1(x)
        if self.batchnorm:
            out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        if self.batchnorm:
            out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        if self.batchnorm:
            out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)
            if self.batchnorm:
                residual = self.bn4(residual)

        out += residual
        out = self.relu(out)

        return out


class SeqEncoder(nn.Module):
    def __init__(self, x_dim: int, x_len: int, h_dim: int, z_dim: int, bias: bool,
                 num_layers: int, batch_norm: bool):
        super().__init__()

        self.bias = bias
        self.batchnorm = batch_norm
        enc = [self._make_layer(x_dim, h_dim, (3, 1, 1))]
        in_dim = h_dim
        window_size = x_len
        for i in range(num_layers - 2):
            out_dim = h_dim * 2 ** i
            enc.append(self._make_layer(in_dim, out_dim, (3, 2, 1)))
            in_dim = out_dim
            window_size = math.floor((window_size + 2 - 3) / 2) + 1

        self.enc = nn.Sequential(*enc)
        self.final_layer = nn.Conv1d(in_dim, z_dim, int(window_size), 1, 0)

    def _make_layer(self, in_dim: int, out_dim: int, conv_param=None):
        downsample = None
        if conv_param is not None:
            downsample = nn.Conv1d(in_dim, out_dim, conv_param[0], conv_param[1], conv_param[2], bias=self.bias)
        elif in_dim != out_dim:
            downsample = nn.Conv1d(in_dim, out_dim, 1, 1, 0, bias=self.bias)

        return ResBlock(in_dim, out_dim, conv_param, downsample=downsample,
                        batchnorm=self.batchnorm, bias=self.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        z = self.enc(x)
        z = self.final_layer(z)
        return z.squeeze(-1)


def make_seq_nets(x_dim: int, config: dict):
    enc_nlayers = config['enc_nlayers']
    enc_hdim = config['enc_hdim']
    z_dim = config['latent_dim']
    x_len = config['x_length']
    num_trans = config['num_trans']
    batch_norm = config['batch_norm']

    enc = nn.ModuleList([
        SeqEncoder(x_dim, x_len, enc_hdim, z_dim, config['enc_bias'], enc_nlayers, batch_norm)
        for _ in range(num_trans + 1)
    ])
    return enc


class NeutralAD(BaseModel):
    def __init__(self, ts_channels: int, seq_len: int, num_trans: int = 4, trans_type: str = 'residual',
                 enc_hdim: int = 32, enc_nlayers: int = 4, trans_nlayers: int = 4, latent_dim: int = 32,
                 batch_norm: bool = False, enc_bias: bool = False,
                 transform_families: Optional[Sequence[Sequence[Transform]]] = None,
                 use_invertible_transforms: bool = False,
                 use_group_transforms: bool = False,
                 use_freq_ortho_transforms: bool = False,
                 keep_base_transforms: bool = True,
                 invertible_cfg: Optional[Dict[str, object]] = None,
                 group_cfg: Optional[Dict[str, object]] = None,
                 freq_cfg: Optional[Dict[str, object]] = None):
        super().__init__()

        self.trans_type = trans_type
        self.z_dim = latent_dim
        transforms: List[Transform] = []
        if keep_base_transforms:
            if num_trans < 1:
                raise ValueError('num_trans must be >= 1 when keep_base_transforms is True')
            transforms.extend(
                [SeqTransformNet(ts_channels, ts_channels, trans_nlayers) for _ in range(num_trans)]
            )
        if use_invertible_transforms:
            invertible_cfg_data = invertible_cfg or default_invertible_cfg()
            transforms.extend(make_invertible_family(ts_channels, invertible_cfg_data))
        if use_group_transforms:
            group_cfg_data = group_cfg or default_group_cfg()
            transforms.extend(make_group_family(ts_channels, group_cfg_data))
        if use_freq_ortho_transforms:
            freq_cfg_data = freq_cfg or default_freq_cfg()
            transforms.extend(make_freq_family(ts_channels, freq_cfg_data))
        if transform_families:
            for family in transform_families:
                transforms.extend(family)
        if not transforms:
            raise ValueError('At least one transform must be configured for NeutralAD')
        self.transform_bank = TransformBank(transforms)
        self.num_trans = len(self.transform_bank)
        config = dict(
            enc_nlayers=enc_nlayers,
            enc_hdim=enc_hdim,
            latent_dim=latent_dim,
            x_length=seq_len,
            num_trans=self.num_trans,
            batch_norm=batch_norm,
            enc_bias=enc_bias
        )
        self.enc = make_seq_nets(ts_channels, config)

    def forward(self, inputs: Tuple[torch.Tensor, ...]) -> torch.Tensor:
        x, = inputs
        x = x.float()
        x = x.permute(0, 2, 1)

        masks = self.transform_bank(x)
        if self.trans_type == 'forward':
            x_t = masks
        elif self.trans_type == 'mul':
            x_t = torch.sigmoid(masks) * x.unsqueeze(1)
        elif self.trans_type == 'residual':
            x_t = masks + x.unsqueeze(1)
        else:
            raise ValueError(f'Unknown trans_type: {self.trans_type}')

        x_cat = torch.cat([x.unsqueeze(1), x_t], 1)
        zs = self.enc[0](x_cat.reshape(-1, x.shape[1], x.shape[2]))
        zs = zs.reshape(x.shape[0], self.num_trans + 1, self.z_dim)

        return zs


def _dcl_score(z: torch.Tensor, temperature: float, eval: bool) -> torch.Tensor:
    z = F.normalize(z, p=2, dim=-1)
    z_ori = z[:, 0]
    z_trans = z[:, 1:]
    num_trans = z.shape[1]

    logits = torch.matmul(z, z.mT) / temperature
    diag_mask = torch.eye(num_trans, device=z.device, dtype=torch.bool)
    logits = logits.masked_fill(diag_mask, float('-inf'))
    trans_logsumexp = torch.logsumexp(logits[:, 1:], dim=-1)

    pos_log = torch.sum(z_trans * z_ori.unsqueeze(1), -1) / temperature
    k_trans = num_trans - 1
    scale = 1 / abs(k_trans * math.log(1.0 / k_trans))

    loss_tensor = (trans_logsumexp - pos_log) * scale

    score = loss_tensor.sum(1)
    if eval:
        return score
    return score.mean()


def _eucdcl_score(z: torch.Tensor, temperature: float, eval: bool) -> torch.Tensor:
    num_trans = z.shape[1]
    logits = -torch.cdist(z, z) / temperature
    diag_mask = torch.eye(num_trans, device=z.device, dtype=torch.bool)
    logits = logits.masked_fill(diag_mask, float('-inf'))
    trans_logsumexp = torch.logsumexp(logits[:, 1:], dim=-1)
    pos_log = logits[:, 1:, 0]

    k_trans = num_trans - 1
    scale = 1 / abs(k_trans * math.log(1.0 / k_trans))
    score = (-pos_log + trans_logsumexp) * scale
    score = score.sum(1)
    if eval:
        return score
    return score.mean()


class NeutralADLoss(Loss):
    def __init__(self, temperature: float = 0.1, use_euclidean: bool = False):
        super().__init__()
        self.temperature = temperature
        self.use_euclidean = use_euclidean

    def forward(self, predictions: Tuple[torch.Tensor, ...], targets: Tuple[torch.Tensor, ...] = None,
                eval: bool = False, *args, **kwargs) -> torch.Tensor:
        z, = predictions
        if self.use_euclidean:
            return _eucdcl_score(z, self.temperature, eval=eval)
        return _dcl_score(z, self.temperature, eval=eval)


class NeutralADAnomalyDetector(AnomalyDetector):
    def __init__(self, model: NeutralAD, loss: NeutralADLoss):
        super().__init__()
        self.model = model
        self.loss = loss

    def compute_online_anomaly_score(self, inputs: Tuple[torch.Tensor, ...]) -> torch.Tensor:
        with torch.inference_mode():
            z = pack_tuple(self.model(inputs))

        return self.loss(z, eval=True)

    def compute_offline_anomaly_score(self, inputs: Tuple[torch.Tensor, ...]) -> torch.Tensor:
        raise NotImplementedError

    def fit(self, dataset: torch.utils.data.DataLoader, **kwargs) -> None:
        pass

    def format_online_targets(self, targets: Tuple[torch.Tensor, ...]) -> torch.Tensor:
        label, = targets
        return label[:, -1]