diff --git a/anomalib/core/callbacks/min_max_normalization.py b/anomalib/core/callbacks/min_max_normalization.py index 3bcef4dc46..e3406d3be9 100644 --- a/anomalib/core/callbacks/min_max_normalization.py +++ b/anomalib/core/callbacks/min_max_normalization.py @@ -73,11 +73,11 @@ def on_predict_batch_end( @staticmethod def _normalize_batch(outputs, pl_module): """Normalize a batch of predictions.""" - stats = pl_module.min_max + stats = pl_module.min_max.cpu() outputs["pred_scores"] = normalize( - outputs["pred_scores"], pl_module.image_threshold.value, stats.min, stats.max + outputs["pred_scores"], pl_module.image_threshold.value.cpu(), stats.min, stats.max ) if "anomaly_maps" in outputs.keys(): outputs["anomaly_maps"] = normalize( - outputs["anomaly_maps"], pl_module.pixel_threshold.value, stats.min, stats.max + outputs["anomaly_maps"], pl_module.pixel_threshold.value.cpu(), stats.min, stats.max ) diff --git a/anomalib/core/model/anomaly_module.py b/anomalib/core/model/anomaly_module.py index 0709bc424a..217c865375 100644 --- a/anomalib/core/model/anomaly_module.py +++ b/anomalib/core/model/anomaly_module.py @@ -48,19 +48,19 @@ def __init__(self, params: Union[DictConfig, ListConfig]): self.loss: Tensor self.callbacks: List[Callback] - self.image_threshold = AdaptiveThreshold(self.hparams.model.threshold.image_default) + self.image_threshold = AdaptiveThreshold(self.hparams.model.threshold.image_default).cpu() self.pixel_threshold = AdaptiveThreshold(self.hparams.model.threshold.pixel_default) - self.training_distribution = AnomalyScoreDistribution() - self.min_max = MinMax() + self.training_distribution = AnomalyScoreDistribution().cpu() + self.min_max = MinMax().cpu() self.model: nn.Module # metrics auroc = AUROC(num_classes=1, pos_label=1, compute_on_step=False) f1_score = F1(num_classes=1, compute_on_step=False) - self.image_metrics = MetricCollection([auroc, f1_score], prefix="image_") - self.pixel_metrics = self.image_metrics.clone(prefix="pixel_") + self.image_metrics = MetricCollection([auroc, f1_score], prefix="image_").cpu() + self.pixel_metrics = self.image_metrics.clone(prefix="pixel_").cpu() def forward(self, batch): # pylint: disable=arguments-differ """Forward-pass input tensor to the module. @@ -111,11 +111,13 @@ def test_step(self, batch, _): # pylint: disable=arguments-differ def validation_step_end(self, val_step_outputs): # pylint: disable=arguments-differ """Called at the end of each validation step.""" + self._outputs_to_cpu(val_step_outputs) self._post_process(val_step_outputs) return val_step_outputs def test_step_end(self, test_step_outputs): # pylint: disable=arguments-differ """Called at the end of each test step.""" + self._outputs_to_cpu(test_step_outputs) self._post_process(test_step_outputs) return test_step_outputs @@ -152,8 +154,10 @@ def _compute_adaptive_threshold(self, outputs): def _collect_outputs(self, image_metric, pixel_metric, outputs): for output in outputs: + image_metric.cpu() image_metric.update(output["pred_scores"], output["label"].int()) if "mask" in output.keys() and "anomaly_maps" in output.keys(): + pixel_metric.cpu() pixel_metric.update(output["anomaly_maps"].flatten(), output["mask"].flatten().int()) def _post_process(self, outputs): @@ -163,6 +167,12 @@ def _post_process(self, outputs): outputs["anomaly_maps"].reshape(outputs["anomaly_maps"].shape[0], -1).max(dim=1).values ) + def _outputs_to_cpu(self, output): + # for output in outputs: + for key, value in output.items(): + if isinstance(value, Tensor): + output[key] = value.cpu() + def _log_metrics(self): """Log computed performance metrics.""" self.log_dict(self.image_metrics) diff --git a/anomalib/core/model/feature_extractor.py b/anomalib/core/model/feature_extractor.py index 2729e65bb8..f616d23280 100644 --- a/anomalib/core/model/feature_extractor.py +++ b/anomalib/core/model/feature_extractor.py @@ -50,10 +50,17 @@ def __init__(self, backbone: nn.Module, layers: Iterable[str]): self.backbone = backbone self.layers = layers self._features = {layer: torch.empty(0) for layer in self.layers} + self.out_dims = [] for layer_id in layers: layer = dict([*self.backbone.named_modules()])[layer_id] layer.register_forward_hook(self.get_features(layer_id)) + # get output dimension of features if available + layer_modules = [*layer.modules()] + for idx in reversed(range(len(layer_modules))): + if hasattr(layer_modules[idx], "out_channels"): + self.out_dims.append(layer_modules[idx].out_channels) + break def get_features(self, layer_id: str) -> Callable: """Get layer features. diff --git a/anomalib/models/__init__.py b/anomalib/models/__init__.py index 0b01bb8c3a..f4b611ebe6 100644 --- a/anomalib/models/__init__.py +++ b/anomalib/models/__init__.py @@ -47,7 +47,7 @@ def get_model(config: Union[DictConfig, ListConfig]) -> AnomalyModule: AnomalyModule: Anomaly Model """ openvino_model_list: List[str] = ["stfpm"] - torch_model_list: List[str] = ["padim", "stfpm", "dfkde", "dfm", "patchcore"] + torch_model_list: List[str] = ["padim", "stfpm", "dfkde", "dfm", "patchcore", "cflow"] model: AnomalyModule if config.openvino: diff --git a/anomalib/models/cflow/README.md b/anomalib/models/cflow/README.md new file mode 100644 index 0000000000..a7cf23ccf2 --- /dev/null +++ b/anomalib/models/cflow/README.md @@ -0,0 +1,3 @@ +# Real-Time Unsupervised Anomaly Detection via Conditional Normalizing Flows + +This is the implementation of the [CW-AD](https://arxiv.org/pdf/2107.12571v1.pdf) paper. diff --git a/anomalib/models/cflow/__init__.py b/anomalib/models/cflow/__init__.py new file mode 100644 index 0000000000..1a4c68fe91 --- /dev/null +++ b/anomalib/models/cflow/__init__.py @@ -0,0 +1,18 @@ +"""Real-Time Unsupervised Anomaly Detection via Conditional Normalizing Flows. + +[CW-AD](https://arxiv.org/pdf/2107.12571v1.pdf) +""" + +# Copyright (C) 2020 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions +# and limitations under the License. diff --git a/anomalib/models/cflow/backbone.py b/anomalib/models/cflow/backbone.py new file mode 100644 index 0000000000..61650e7319 --- /dev/null +++ b/anomalib/models/cflow/backbone.py @@ -0,0 +1,100 @@ +"""Helper functions to create backbone model.""" + +# Copyright (C) 2020 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions +# and limitations under the License. + +import math + +import FrEIA.framework as Ff +import FrEIA.modules as Fm +import torch +from FrEIA.framework.sequence_inn import SequenceINN +from torch import nn + + +def positional_encoding_2d(condition_vector: int, height: int, width: int) -> torch.Tensor: + """Creates embedding to store relative position of the feature vector using sine and cosine functions. + + Args: + condition_vector (int): Length of the condition vector + height (int): H of the positions + width (int): W of the positions + + Raises: + ValueError: Cannot generate encoding with conditional vector length not as multiple of 4 + + Returns: + torch.Tensor: condition_vector x HEIGHT x WIDTH position matrix + """ + if condition_vector % 4 != 0: + raise ValueError(f"Cannot use sin/cos positional encoding with odd dimension (got dim={condition_vector})") + pos_encoding = torch.zeros(condition_vector, height, width) + # Each dimension use half of condition_vector + condition_vector = condition_vector // 2 + div_term = torch.exp(torch.arange(0.0, condition_vector, 2) * -(math.log(1e4) / condition_vector)) + pos_w = torch.arange(0.0, width).unsqueeze(1) + pos_h = torch.arange(0.0, height).unsqueeze(1) + pos_encoding[0:condition_vector:2, :, :] = ( + torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1) + ) + pos_encoding[1:condition_vector:2, :, :] = ( + torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1) + ) + pos_encoding[condition_vector::2, :, :] = ( + torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width) + ) + pos_encoding[condition_vector + 1 :: 2, :, :] = ( + torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width) + ) + return pos_encoding + + +def subnet_fc(dims_in: int, dims_out: int): + """Subnetwork which predicts the affine coefficients. + + Args: + dims_in (int): input dimensions + dims_out (int): output dimensions + + Returns: + nn.Sequential: Feed-forward subnetwork + """ + return nn.Sequential(nn.Linear(dims_in, 2 * dims_in), nn.ReLU(), nn.Linear(2 * dims_in, dims_out)) + + +def cflow_head(condition_vector: int, coupling_blocks: int, clamp_alpha: float, n_features: int) -> SequenceINN: + """Create invertible decoder network. + + Args: + condition_vector (int): length of the condition vector + coupling_blocks (int): number of coupling blocks to build the decoder + clamp_alpha (float): clamping value to avoid exploding values + n_features (int): number of decoder features + + Returns: + SequenceINN: decoder network block + """ + coder = Ff.SequenceINN(n_features) + print("CNF coder:", n_features) + for _ in range(coupling_blocks): + coder.append( + Fm.AllInOneBlock, + cond=0, + cond_shape=(condition_vector,), + subnet_constructor=subnet_fc, + affine_clamping=clamp_alpha, + global_affine_type="SOFTPLUS", + permute_soft=True, + ) + return coder diff --git a/anomalib/models/cflow/config.yaml b/anomalib/models/cflow/config.yaml new file mode 100644 index 0000000000..0a17079f2c --- /dev/null +++ b/anomalib/models/cflow/config.yaml @@ -0,0 +1,96 @@ +dataset: + name: mvtec + format: mvtec + path: ./datasets/MVTec + url: ftp://guest:GU.205dldo@ftp.softronics.ch/mvtec_anomaly_detection/mvtec_anomaly_detection.tar.xz + category: leather + task: segmentation + label_format: None + image_size: 256 + train_batch_size: 16 + test_batch_size: 16 + inference_batch_size: 16 + fiber_batch_size: 64 + num_workers: 36 + +model: + name: cflow + backbone: resnet18 + layers: + - layer2 + - layer3 + - layer4 + decoder: freia-cflow + condition_vector: 128 + coupling_blocks: 8 + clamp_alpha: 1.9 + lr: 0.0001 + early_stopping: + patience: 3 + metric: pixel_AUROC + mode: max + normalization_method: min_max # options: [null, min_max, cdf] + threshold: + image_default: 0 + pixel_default: 0 + adaptive: true + +project: + seed: 0 + path: ./results + log_images_to: [local] + logger: false + save_to_csv: false + +# PL Trainer Args. Don't add extra parameter here. +trainer: + accelerator: null + accumulate_grad_batches: 1 + amp_backend: native + amp_level: O2 + auto_lr_find: false + auto_scale_batch_size: false + auto_select_gpus: false + benchmark: false + check_val_every_n_epoch: 1 + checkpoint_callback: true + default_root_dir: null + deterministic: true + distributed_backend: null + fast_dev_run: false + flush_logs_every_n_steps: 100 + gpus: 1 + gradient_clip_val: 0 + limit_predict_batches: 1.0 + limit_test_batches: 1.0 + limit_train_batches: 1.0 + limit_val_batches: 1.0 + log_every_n_steps: 50 + log_gpu_memory: null + max_epochs: 50 + max_steps: null + min_epochs: null + min_steps: null + move_metrics_to_cpu: false + multiple_trainloader_mode: max_size_cycle + num_nodes: 1 + num_processes: 1 + num_sanity_val_steps: 0 + overfit_batches: 0.0 + plugins: null + precision: 32 + prepare_data_per_node: true + process_position: 0 + profiler: null + progress_bar_refresh_rate: null + reload_dataloaders_every_epoch: false + replace_sampler_ddp: true + stochastic_weight_avg: false + sync_batchnorm: false + terminate_on_nan: false + tpu_cores: null + track_grad_norm: -1 + truncated_bptt_steps: null + val_check_interval: 1.0 + weights_save_path: null + weights_summary: top diff --git a/anomalib/models/cflow/model.py b/anomalib/models/cflow/model.py new file mode 100644 index 0000000000..d9385eaaf8 --- /dev/null +++ b/anomalib/models/cflow/model.py @@ -0,0 +1,344 @@ +"""CFLOW: Real-Time Unsupervised Anomaly Detection via Conditional Normalizing Flows. + +https://arxiv.org/pdf/2107.12571v1.pdf +""" + +# Copyright (C) 2020 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions +# and limitations under the License. + +from typing import List, Tuple, Union, cast + +import einops +import numpy as np +import torch +import torch.nn.functional as F +import torchvision +from omegaconf import DictConfig, ListConfig +from pytorch_lightning.callbacks import EarlyStopping +from torch import Tensor, nn, optim + +from anomalib.core.model import AnomalyModule +from anomalib.core.model.feature_extractor import FeatureExtractor +from anomalib.models.cflow.backbone import cflow_head, positional_encoding_2d + +__all__ = ["AnomalyMapGenerator", "CflowModel", "CflowLightning"] + + +def get_logp(dim_feature_vector: int, p_u: torch.Tensor, logdet_j: torch.Tensor) -> torch.Tensor: + """Returns the log likelihood estimation. + + Args: + dim_feature_vector (int): Dimensions of the condition vector + p_u (torch.Tensor): Random variable u + logdet_j (torch.Tensor): log of determinant of jacobian returned from the invertable decoder + + Returns: + torch.Tensor: Log probability + """ + ln_sqrt_2pi = -np.log(np.sqrt(2 * np.pi)) # ln(sqrt(2*pi)) + logp = dim_feature_vector * ln_sqrt_2pi - 0.5 * torch.sum(p_u ** 2, 1) + logdet_j + return logp + + +class AnomalyMapGenerator: + """Generate Anomaly Heatmap.""" + + def __init__( + self, + image_size: Union[ListConfig, Tuple], + pool_layers: List[str], + ): + self.distance = torch.nn.PairwiseDistance(p=2, keepdim=True) + self.image_size = image_size if isinstance(image_size, tuple) else tuple(image_size) + self.pool_layers: List[str] = pool_layers + + def compute_anomaly_map( + self, distribution: Union[List[Tensor], List[List]], height: List[int], width: List[int] + ) -> Tensor: + """Compute the layer map based on likelihood estimation. + + Args: + distribution: Probability distribution for each decoder block + height: blocks height + width: blocks width + + Returns: + Final Anomaly Map + + """ + + test_map: List[Tensor] = [] + for layer_idx in range(len(self.pool_layers)): + test_norm = torch.tensor(distribution[layer_idx], dtype=torch.double) # pylint: disable=not-callable + test_norm -= torch.max(test_norm) # normalize likelihoods to (-Inf:0] by subtracting a constant + test_prob = torch.exp(test_norm) # convert to probs in range [0:1] + test_mask = test_prob.reshape(-1, height[layer_idx], width[layer_idx]) + # upsample + test_map.append( + F.interpolate( + test_mask.unsqueeze(1), size=self.image_size, mode="bilinear", align_corners=True + ).squeeze() + ) + # score aggregation + score_map = torch.zeros_like(test_map[0]) + for layer_idx in range(len(self.pool_layers)): + score_map += test_map[layer_idx] + score_mask = score_map + # invert probs to anomaly scores + anomaly_map = score_mask.max() - score_mask + + return anomaly_map + + def __call__(self, **kwargs: Union[List[Tensor], List[int], List[List]]) -> Tensor: + """Returns anomaly_map. + + Expects `distribution`, `height` and 'width' keywords to be passed explicitly + + Example + >>> anomaly_map_generator = AnomalyMapGenerator(image_size=tuple(hparams.model.input_size), + >>> pool_layers=pool_layers) + >>> output = self.anomaly_map_generator(distribution=dist, height=height, width=width) + + Raises: + ValueError: `distribution`, `height` and 'width' keys are not found + + Returns: + torch.Tensor: anomaly map + """ + if not ("distribution" in kwargs and "height" in kwargs and "width" in kwargs): + raise KeyError(f"Expected keys `distribution`, `height` and `width`. Found {kwargs.keys()}") + + # placate mypy + distribution: List[Tensor] = cast(List[Tensor], kwargs["distribution"]) + height: List[int] = cast(List[int], kwargs["height"]) + width: List[int] = cast(List[int], kwargs["width"]) + return self.compute_anomaly_map(distribution, height, width) + + +class CflowModel(nn.Module): + """CFLOW: Conditional Normalizing Flows.""" + + def __init__(self, hparams: Union[DictConfig, ListConfig]): + super().__init__() + + self.backbone = getattr(torchvision.models, hparams.model.backbone) + self.fiber_batch_size = hparams.dataset.fiber_batch_size + self.condition_vector: int = hparams.model.condition_vector + self.dec_arch = hparams.model.decoder + self.pool_layers = hparams.model.layers + + self.encoder = FeatureExtractor(backbone=self.backbone(pretrained=True), layers=self.pool_layers) + self.pool_dims = self.encoder.out_dims + self.decoders = nn.ModuleList( + [ + cflow_head(self.condition_vector, hparams.model.coupling_blocks, hparams.model.clamp_alpha, pool_dim) + for pool_dim in self.pool_dims + ] + ) + + # encoder model is fixed + for parameters in self.encoder.parameters(): + parameters.requires_grad = False + + self.anomaly_map_generator = AnomalyMapGenerator( + image_size=tuple(hparams.model.input_size), pool_layers=self.pool_layers + ) + + def forward(self, images): + """Forward-pass images into the network to extract encoder features and compute probability. + + Args: + images: Batch of images. + + Returns: + Predicted anomaly maps. + + """ + + activation = self.encoder(images) + + distribution = [[] for _ in self.pool_layers] + + height: List[int] = [] + width: List[int] = [] + for layer_idx, layer in enumerate(self.pool_layers): + encoder_activations = activation[layer].detach() # BxCxHxW + + batch_size, dim_feature_vector, im_height, im_width = encoder_activations.size() + image_size = im_height * im_width + embedding_length = batch_size * image_size # number of rows in the conditional vector + + height.append(im_height) + width.append(im_width) + # repeats positional encoding for the entire batch 1 C H W to B C H W + pos_encoding = einops.repeat( + positional_encoding_2d(self.condition_vector, im_height, im_width).unsqueeze(0), + "b c h w-> (tile b) c h w", + tile=batch_size, + ).to(images.device) + c_r = einops.rearrange(pos_encoding, "b c h w -> (b h w) c") # BHWxP + e_r = einops.rearrange(encoder_activations, "b c h w -> (b h w) c") # BHWxC + decoder = self.decoders[layer_idx].to(images.device) + + # Sometimes during validation, the last batch E / N is not a whole number. Hence we need to add 1. + # It is assumed that during training that E / N is a whole number as no errors were discovered during + # testing. In case it is observed in the future, we can use only this line and ensure that FIB is at + # least 1 or set `drop_last` in the dataloader to drop the last non-full batch. + fiber_batches = embedding_length // self.fiber_batch_size + int( + embedding_length % self.fiber_batch_size > 0 + ) + + for batch_num in range(fiber_batches): # per-fiber processing + if batch_num < (fiber_batches - 1): + idx = torch.arange(batch_num * self.fiber_batch_size, (batch_num + 1) * self.fiber_batch_size) + else: # When non-full batch is encountered batch_num+1 * N will go out of bounds + idx = torch.arange(batch_num * self.fiber_batch_size, embedding_length) + c_p = c_r[idx] # NxP + e_p = e_r[idx] # NxC + # decoder returns the transformed variable z and the log Jacobian determinant + p_u, log_jac_det = decoder(e_p, [c_p]) + # + decoder_log_prob = get_logp(dim_feature_vector, p_u, log_jac_det) + log_prob = decoder_log_prob / dim_feature_vector # likelihood per dim + distribution[layer_idx] = distribution[layer_idx] + log_prob.detach().tolist() + + output = self.anomaly_map_generator(distribution=distribution, height=height, width=width) + + return output.to(images.device) + + +class CflowLightning(AnomalyModule): + """PL Lightning Module for the CFLOW algorithm.""" + + def __init__(self, hparams): + super().__init__(hparams) + + self.model: CflowModel = CflowModel(hparams) + self.loss_val = 0 + self.automatic_optimization = False + + def configure_callbacks(self): + """Configure model-specific callbacks.""" + early_stopping = EarlyStopping( + monitor=self.hparams.model.early_stopping.metric, + patience=self.hparams.model.early_stopping.patience, + mode=self.hparams.model.early_stopping.mode, + ) + return [early_stopping] + + def configure_optimizers(self) -> torch.optim.Optimizer: + """Configures optimizers for each decoder. + + Returns: + Optimizer: Adam optimizer for each decoder + """ + decoders_parameters = [] + for decoder_idx in range(len(self.model.pool_layers)): + decoders_parameters.extend(list(self.model.decoders[decoder_idx].parameters())) + + optimizer = optim.Adam( + params=decoders_parameters, + lr=self.hparams.model.lr, + ) + return optimizer + + def training_step(self, batch, _): # pylint: disable=arguments-differ + """Training Step of CFLOW. + + For each batch, decoder layers are trained with a dynamic fiber batch size. + Training step is performed manually as multiple training steps are involved + per batch of input images + + Args: + batch: Input batch + _: Index of the batch. + + Returns: + Loss value for the batch + + """ + opt = self.optimizers() + self.model.encoder.eval() + + images = batch["image"] + activation = self.model.encoder(images) + avg_loss = torch.zeros([1], dtype=torch.float64).to(images.device) + + height = [] + width = [] + for layer_idx, layer in enumerate(self.model.pool_layers): + encoder_activations = activation[layer].detach() # BxCxHxW + + batch_size, dim_feature_vector, im_height, im_width = encoder_activations.size() + image_size = im_height * im_width + embedding_length = batch_size * image_size # number of rows in the conditional vector + + height.append(im_height) + width.append(im_width) + # repeats positional encoding for the entire batch 1 C H W to B C H W + pos_encoding = einops.repeat( + positional_encoding_2d(self.model.condition_vector, im_height, im_width).unsqueeze(0), + "b c h w-> (tile b) c h w", + tile=batch_size, + ).to(images.device) + c_r = einops.rearrange(pos_encoding, "b c h w -> (b h w) c") # BHWxP + e_r = einops.rearrange(encoder_activations, "b c h w -> (b h w) c") # BHWxC + perm = torch.randperm(embedding_length) # BHW + decoder = self.model.decoders[layer_idx].to(images.device) + + fiber_batches = embedding_length // self.model.fiber_batch_size # number of fiber batches + assert fiber_batches > 0, "Make sure we have enough fibers, otherwise decrease N or batch-size!" + + for batch_num in range(fiber_batches): # per-fiber processing + opt.zero_grad() + if batch_num < (fiber_batches - 1): + idx = torch.arange( + batch_num * self.model.fiber_batch_size, (batch_num + 1) * self.model.fiber_batch_size + ) + else: # When non-full batch is encountered batch_num * N will go out of bounds + idx = torch.arange(batch_num * self.model.fiber_batch_size, embedding_length) + # get random vectors + c_p = c_r[perm[idx]] # NxP + e_p = e_r[perm[idx]] # NxC + # decoder returns the transformed variable z and the log Jacobian determinant + p_u, log_jac_det = decoder(e_p, [c_p]) + # + decoder_log_prob = get_logp(dim_feature_vector, p_u, log_jac_det) + log_prob = decoder_log_prob / dim_feature_vector # likelihood per dim + loss = -F.logsigmoid(log_prob) + self.manual_backward(loss.mean()) + opt.step() + avg_loss += loss.sum() + + return {"loss": avg_loss} + + def validation_step(self, batch, _): # pylint: disable=arguments-differ + """Validation Step of CFLOW. + + Similar to the training step, encoder features + are extracted from the CNN for each batch, and anomaly + map is computed. + + Args: + batch: Input batch + _: Index of the batch. + + Returns: + Dictionary containing images, anomaly maps, true labels and masks. + These are required in `validation_epoch_end` for feature concatenation. + + """ + batch["anomaly_maps"] = self.model(batch["image"]) + + return batch diff --git a/requirements/base.txt b/requirements/base.txt index 623c3bb527..3d17b910f8 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -1,6 +1,8 @@ albumentations==1.1.0 attrdict==2.0.1 defusedxml==0.7.1 +einops==0.3.2 +FrEIA @ git+https://github.com/VLL-HD/FrEIA.git kornia==0.5.6 lxml==4.6.5 matplotlib==3.4.3 diff --git a/tests/core/callbacks/normalization_callback/test_normalization_callback.py b/tests/core/callbacks/normalization_callback/test_normalization_callback.py index e6262999bb..14237edfda 100644 --- a/tests/core/callbacks/normalization_callback/test_normalization_callback.py +++ b/tests/core/callbacks/normalization_callback/test_normalization_callback.py @@ -39,5 +39,7 @@ def test_normalizer(): # performance should be the same for metric in ["image_AUROC", "image_F1"]: - assert results_without_normalization[0][metric] == results_with_cdf_normalization[0][metric] - assert results_without_normalization[0][metric] == results_with_minmax_normalization[0][metric] + assert round(results_without_normalization[0][metric], 3) == round(results_with_cdf_normalization[0][metric], 3) + assert round(results_without_normalization[0][metric], 3) == round( + results_with_minmax_normalization[0][metric], 3 + ) diff --git a/tests/models/test_model.py b/tests/models/test_model.py index 573ec9ea41..9b11c57455 100644 --- a/tests/models/test_model.py +++ b/tests/models/test_model.py @@ -155,6 +155,7 @@ def _test_model_load(self, config, datamodule, results): ("stfpm", False), ("stfpm", True), ("patchcore", False), + ("cflow", False), ], ) @pytest.mark.flaky(max_runs=3)