diff --git a/bdpy/dl/torch/dataset.py b/bdpy/dl/torch/dataset.py new file mode 100644 index 00000000..c821592f --- /dev/null +++ b/bdpy/dl/torch/dataset.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +from typing import Iterable, Callable, Dict + +from pathlib import Path + +from PIL import Image +import numpy as np +from torch.utils.data import Dataset + +from bdpy.dataform import DecodedFeatures, Features + + +_FeatureTypeNP = Dict[str, np.ndarray] + + +def _removesuffix(s: str, suffix: str) -> str: + """Remove suffix from string. + + Note + ---- + This function is available from Python 3.9 as `str.removesuffix`. We can + remove this function when we drop support for Python 3.8. + + Parameters + ---------- + s : str + String. + suffix : str + Suffix to remove. + + Returns + ------- + str + String without suffix. + """ + if suffix and s.endswith(suffix): + return s[: -len(suffix)] + return s[:] + + +class FeaturesDataset(Dataset): + """Dataset of features. + + Parameters + ---------- + root_path : str | Path + Path to the root directory of features. + layer_path_names : Iterable[str] + List of layer path names. Each layer path name is used to get features + from the root directory so that the layer path name must be a part of + the path to the layer. + stimulus_names : list[str], optional + List of stimulus names. If None, all stimulus names are used. + transform : callable, optional + Callable object which is used to transform features. The callable object + must take a dict of features and return a dict of features. + """ + + def __init__( + self, + root_path: str | Path, + layer_path_names: Iterable[str], + stimulus_names: list[str] | None = None, + transform: Callable[[_FeatureTypeNP], _FeatureTypeNP] | None = None, + ): + self._features_store = Features(Path(root_path).as_posix()) + self._layer_path_names = layer_path_names + if stimulus_names is None: + stimulus_names = self._features_store.labels + self._stimulus_names = stimulus_names + self._transform = transform + + def __len__(self) -> int: + return len(self._stimulus_names) + + def __getitem__(self, index: int) -> _FeatureTypeNP: + stimulus_name = self._stimulus_names[index] + features = {} + for layer_path_name in self._layer_path_names: + feature = self._features_store.get( + layer=layer_path_name, label=stimulus_name + ) + feature = feature[0] # NOTE: remove batch axis + features[layer_path_name] = feature + if self._transform is not None: + features = self._transform(features) + return features + + +class DecodedFeaturesDataset(Dataset): + """Dataset of decoded features. + + Parameters + ---------- + root_path : str | Path + Path to the root directory of decoded features. + layer_path_names : Iterable[str] + List of layer path names. Each layer path name is used to get features + from the root directory so that the layer path name must be a part of + the path to the layer. + subject_id : str + ID of the subject. + roi : str + ROI name. + stimulus_names : list[str], optional + List of stimulus names. If None, all stimulus names are used. + transform : callable, optional + Callable object which is used to transform features. The callable object + must take a dict of features and return a dict of features. + """ + + def __init__( + self, + root_path: str | Path, + layer_path_names: Iterable[str], + subject_id: str, + roi: str, + stimulus_names: list[str] | None = None, + transform: Callable[[_FeatureTypeNP], _FeatureTypeNP] | None = None, + ): + self._decoded_features_store = DecodedFeatures(Path(root_path).as_posix()) + self._layer_path_names = layer_path_names + self._subject_id = subject_id + self._roi = roi + if stimulus_names is None: + stimulus_names = self._decoded_features_store.labels + assert stimulus_names is not None + self._stimulus_names = stimulus_names + self._transform = transform + + def __len__(self) -> int: + return len(self._stimulus_names) + + def __getitem__(self, index: int) -> _FeatureTypeNP: + stimulus_name = self._stimulus_names[index] + decoded_features = {} + for layer_path_name in self._layer_path_names: + decoded_feature = self._decoded_features_store.get( + layer=layer_path_name, + label=stimulus_name, + subject=self._subject_id, + roi=self._roi, + ) + decoded_feature = decoded_feature[0] # NOTE: remove batch axis + decoded_features[layer_path_name] = decoded_feature + if self._transform is not None: + decoded_features = self._transform(decoded_features) + return decoded_features + + +class ImageDataset(Dataset): + """Dataset of images. + + Parameters + ---------- + root_path : str | Path + Path to the root directory of images. + stimulus_names : list[str], optional + List of stimulus names. If None, all stimulus names are used. + extension : str, optional + Extension of the image files. + """ + + def __init__( + self, + root_path: str | Path, + stimulus_names: list[str] | None = None, + extension: str = "jpg", + ): + self.root_path = root_path + if stimulus_names is None: + stimulus_names = [ + _removesuffix(path.name, "." + extension) + for path in Path(root_path).glob(f"*{extension}") + ] + self._stimulus_names = stimulus_names + self._extension = extension + + def __len__(self): + return len(self._stimulus_names) + + def __getitem__(self, index: int): + stimulus_name = self._stimulus_names[index] + image = Image.open(Path(self.root_path) / f"{stimulus_name}.{self._extension}") + image = image.convert("RGB") + return np.array(image) / 255.0, stimulus_name + + +class RenameFeatureKeys: + def __init__(self, mapping: dict[str, str]): + self._mapping = mapping + + def __call__(self, features: _FeatureTypeNP) -> _FeatureTypeNP: + return {self._mapping.get(key, key): value for key, value in features.items()} diff --git a/bdpy/dl/torch/domain/__init__.py b/bdpy/dl/torch/domain/__init__.py new file mode 100644 index 00000000..05a08fa1 --- /dev/null +++ b/bdpy/dl/torch/domain/__init__.py @@ -0,0 +1 @@ +from .core import Domain, InternalDomain, IrreversibleDomain, ComposedDomain, KeyValueDomain \ No newline at end of file diff --git a/bdpy/dl/torch/domain/core.py b/bdpy/dl/torch/domain/core.py new file mode 100644 index 00000000..6792cb95 --- /dev/null +++ b/bdpy/dl/torch/domain/core.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Iterable, TypeVar, Generic +import warnings + +import torch.nn as nn + +_T = TypeVar("_T") + + +class Domain(nn.Module, ABC, Generic[_T]): + """Base class for stimulus domain. + + This class is used to convert data between each domain and library's internal common space. + Suppose that we have two functions `f: X -> Y_1` and `g: Y_2 -> Z` and want to compose them. + Here, `X`, `Y_1`, `Y_2`, and `Z` are different domains and assume that `Y_1` and `Y_2` are + the similar domain that can be converted to each other. + Then, we can compose `f` and `g` as `g . t . f(x)`, where `t: Y_1 -> Y_2` is the domain + conversion function. This class is used to implement `t`. + + The subclasses of this class should implement `send` and `receive` methods. The `send` method + converts data from the original domain (`Y_1` or `Y_2`) to the internal common space (`Y_0`), + and the `receive` method converts data from the internal common space to the original domain. + By implementing domain class for `Y_1` and `Y_2`, we can construct the domain conversion function + `t` as `t = Y_2.receive . Y_1.send`. + + Note that the subclasses of this class do not necessarily guarantee the reversibility of `send` + and `receive` methods. If the domain conversion is irreversible, the subclasses should inherit + `IrreversibleDomain` class instead of this class. + """ + + @abstractmethod + def send(self, x: _T) -> _T: + """Send stimulus to the internal common space from each domain. + + Parameters + ---------- + x : _T + Data in the original domain. + + Returns + ------- + _T + Data in the internal common space. + """ + pass + + @abstractmethod + def receive(self, x: _T) -> _T: + """Receive data from the internal common space to each domain. + + Parameters + ---------- + x : _T + Data in the internal common space. + + Returns + ------- + _T + Data in the original domain. + """ + pass + + +class InternalDomain(Domain, Generic[_T]): + """The internal common space. + + The domain class which defines the internal common space. This class + receives and sends data as it is. + """ + + def send(self, x: _T) -> _T: + return x + + def receive(self, x: _T) -> _T: + return x + + +class IrreversibleDomain(Domain, Generic[_T]): + """The domain which cannot be reversed. + + This class is used to convert data between each domain and library's + internal common space. Note that the subclasses of this class do not + guarantee the reversibility of `send` and `receive` methods. + """ + + def __init__(self) -> None: + super().__init__() + warnings.warn( + f"{self.__class__.__name__} is an irreversible domain. " \ + "It does not guarantee the reversibility of `send` and `receive` " \ + "methods. Please use the combination of `send` and `receive` methods " \ + "with caution.", + RuntimeWarning, + ) + + def send(self, x: _T) -> _T: + return x + + def receive(self, x: _T) -> _T: + return x + + +class ComposedDomain(Domain, Generic[_T]): + """The domain composed of multiple sub-domains. + + Suppose we have list of domain objects `domains = [d_0, d_1, ..., d_n]`. + Then, `ComposedDomain(domains)` accesses the data in the original domain `D` + as `d_n.receive . ... d_1.receive . d_0.receive(x)` from the internal common space `D_0`. + + Parameters + ---------- + domains : Iterable[Domain] + Sub-domains to compose. + + Examples + -------- + >>> import numpy as np + >>> import torch + >>> from bdpy.dl.torch.domain import ComposedDomain + >>> from bdpy.dl.torch.domain.image_domain import AffineDomain, BGRDomain + >>> composed_domain = ComposedDomain([ + ... AffineDomain(0.5, 1), + ... BGRDomain(), + ... ]) + >>> image = torch.randn(1, 3, 64, 64).clamp(-0.5, 0.5) + >>> image.shape + torch.Size([1, 3, 64, 64]) + >>> composed_domain.send(image).shape + torch.Size([1, 3, 64, 64]) + >>> print(composed_domain.send(image).min().item(), composed_domain.send(image).max().item()) + 0.0 1.0 + """ + + def __init__(self, domains: Iterable[Domain]) -> None: + super().__init__() + self.domains = nn.ModuleList(domains) + + def send(self, x: _T) -> _T: + for domain in reversed(self.domains): + x = domain.send(x) + return x + + def receive(self, x: _T) -> _T: + for domain in self.domains: + x = domain.receive(x) + return x + + +class KeyValueDomain(Domain, Generic[_T]): + """The domain which converts key-value pairs. + + This class is used to convert key-value pairs between each domain and library's + internal common space. + + Parameters + ---------- + domain_mapper : dict[str, Domain] + Dictionary that maps keys to domains. + """ + + def __init__(self, domain_mapper: dict[str, Domain]) -> None: + super().__init__() + self.domain_mapper = domain_mapper + + def send(self, x: dict[str, _T]) -> dict[str, _T]: + return { + key: self.domain_mapper[key].send(value) for key, value in x.items() + } + + def receive(self, x: dict[str, _T]) -> dict[str, _T]: + return { + key: self.domain_mapper[key].receive(value) for key, value in x.items() + } diff --git a/bdpy/dl/torch/domain/feature_domain.py b/bdpy/dl/torch/domain/feature_domain.py new file mode 100644 index 00000000..ad658384 --- /dev/null +++ b/bdpy/dl/torch/domain/feature_domain.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +from typing import Dict + +import torch + +from .core import Domain + +_FeatureType = Dict[str, torch.Tensor] + + +def _lnd2nld(feature: torch.Tensor) -> torch.Tensor: + """Convert features having the shape of (L, N, D) to (N, L, D).""" + return feature.permute(1, 0, 2) + +def _nld2lnd(feature: torch.Tensor) -> torch.Tensor: + """Convert features having the shape of (N, L, D) to (L, N, D).""" + return feature.permute(1, 0, 2) + + +class ArbitraryFeatureKeyDomain(Domain): + def __init__( + self, + to_internal: dict[str, str] | None = None, + to_self: dict[str, str] | None = None, + ): + super().__init__() + + if to_internal is None and to_self is None: + raise ValueError("Either to_internal or to_self must be specified.") + + if to_internal is None: + to_internal = {value: key for key, value in to_self.items()} + elif to_self is None: + to_self = {value: key for key, value in to_internal.items()} + + self._to_internal = to_internal + self._to_self = to_self + + def send(self, features: _FeatureType) -> _FeatureType: + return {self._to_internal.get(key, key): value for key, value in features.items()} + + def receive(self, features: _FeatureType) -> _FeatureType: + return {self._to_self.get(key, key): value for key, value in features.items()} diff --git a/bdpy/dl/torch/domain/image_domain.py b/bdpy/dl/torch/domain/image_domain.py new file mode 100644 index 00000000..74b49b49 --- /dev/null +++ b/bdpy/dl/torch/domain/image_domain.py @@ -0,0 +1,229 @@ +"""Image domains for PyTorch. + +This module provides image domains for PyTorch. The image domains are used to +convert images between each domain and library's internal common space. +The internal common space is defined as follows: + +- Channel axis: 1 +- Pixel range: [0, 1] +- Image size: arbitrary +- Color space: RGB +""" + +from __future__ import annotations + +import warnings + +import numpy as np +import torch +from torchvision.transforms import InterpolationMode, Resize + +from .core import Domain, InternalDomain, IrreversibleDomain, ComposedDomain + + +def _bgr2rgb(images: torch.Tensor) -> torch.Tensor: + """Convert images from BGR to RGB""" + return images[:, [2, 1, 0], ...] + + +def _rgb2bgr(images: torch.Tensor) -> torch.Tensor: + """Convert images from RGB to BGR""" + return images[:, [2, 1, 0], ...] + + +def _to_channel_first(images: torch.Tensor) -> torch.Tensor: + """Convert images from channel last to channel first""" + return images.permute(0, 3, 1, 2) + + +def _to_channel_last(images: torch.Tensor) -> torch.Tensor: + """Convert images from channel first to channel last""" + return images.permute(0, 2, 3, 1) + + +# NOTE: The internal common space for images is defined as follows: +# - Channel axis: 1 +# - Pixel range: [0, 1] +# - Image size: arbitrary +# - Color space: RGB +Zero2OneImageDomain = InternalDomain[torch.Tensor] + + +class AffineDomain(Domain): + """Image domain shifted by center and scaled by scale. + + This domain is used to convert images in [0, 1] to images in [-center, scale-center]. + In other words, the pixel intensity p in [0, 1] is converted to p * scale - center. + + Parameters + ---------- + center : float | np.ndarray + Center of the affine transformation. + If center.ndim == 0, it must be scalar. + If center.ndim == 1, it must be 1D vector (C,). + If center.ndim == 3, it must be 3D vector (1, C, W, H). + scale : float | np.ndarray + Scale of the affine transformation. + If scale.ndim == 0, it must be scalar. + If scale.ndim == 1, it must be 1D vector (C,). + If scale.ndim == 3, it must be 3D vector (1, C, W, H). + device : torch.device | None + Device to send/receive images. + dtype : torch.dtype | None + Data type to send/receive images. + """ + + def __init__( + self, + center: float | np.ndarray, + scale: float | np.ndarray, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> None: + super().__init__() + + if isinstance(center, (float, int)) or center.ndim == 0: + center = np.array([center])[np.newaxis, np.newaxis, np.newaxis] + elif center.ndim == 1: # 1D vector (C,) + center = center[np.newaxis, :, np.newaxis, np.newaxis] + elif center.ndim == 3: # 3D vector (1, C, W, H) + center = center[np.newaxis] + else: + raise ValueError( + f"center must be 1D or 3D vector, but got {center.ndim}D vector." + ) + if isinstance(scale, (float, int)) or scale.ndim == 0: + scale = np.array([scale])[np.newaxis, np.newaxis, np.newaxis] + elif scale.ndim == 1: # 1D vector (C,) + scale = scale[np.newaxis, :, np.newaxis, np.newaxis] + elif scale.ndim == 3: # 3D vector (1, C, W, H) + scale = scale[np.newaxis] + else: + raise ValueError( + f"scale must be scalar or 1D or 3D vector, but got {scale.ndim}D vector." + ) + + self._center = torch.from_numpy(center).to(device=device, dtype=dtype) + self._scale = torch.from_numpy(scale).to(device=device, dtype=dtype) + + def send(self, images: torch.Tensor) -> torch.Tensor: + return (images + self._center) / self._scale + + def receive(self, images: torch.Tensor) -> torch.Tensor: + return images * self._scale - self._center + + +class BGRDomain(Domain): + """Image domain for BGR images.""" + + def send(self, images: torch.Tensor) -> torch.Tensor: + return _bgr2rgb(images) + + def receive(self, images: torch.Tensor) -> torch.Tensor: + return _rgb2bgr(images) + + +class PILDomainWithExplicitCrop(IrreversibleDomain): + """Image domain for PIL images. + + - Channel axis: 3 + - Pixel range: [0, 255] + - Image size: arbitrary + - Color space: RGB + """ + + def send(self, images: torch.Tensor) -> torch.Tensor: + return _to_channel_first(images) / 255.0 # to [0, 1.0] + + def receive(self, images: torch.Tensor) -> torch.Tensor: + warnings.warn( + "`PILDominWithExplicitCrop.receive` performs explicit cropping. " \ + "It could be affected to the gradient computation. " \ + "Please do not use this domain inside the optimization.", + RuntimeWarning, + ) + + images = _to_channel_last(images) * 255.0 + + # Crop values to [0, 255] + return torch.clamp(images, 0, 255) + + +class BdPyVGGDomain(ComposedDomain): + """Image domain for VGG architecture defined in BdPy. + + - Channel axis: 1 + - Pixel range: + - red: [-123, 132] + - green: [-117, 138] + - blue: [-104, 151] + - Image size: arbitrary + - Color space: BGR + + Parameters + ---------- + device : torch.device | None + Device to send/receive images. + dtype : torch.dtype | None + Data type to send/receive images. + + Notes + ----- + The pixel ranges of this domain are derived from the mean vector of ImageNet ([123, 117, 104]). + """ + + def __init__( + self, *, device: torch.device | None = None, dtype: torch.dtype | None = None + ) -> None: + super().__init__( + [ + AffineDomain( + center=np.array([123.0, 117.0, 104.0]), + scale=255.0, + device=device, + dtype=dtype, + ), + BGRDomain(), + ] + ) + + +class FixedResolutionDomain(IrreversibleDomain): + """Image domain for images with fixed resolution. + + Parameters + ---------- + image_shape : tuple[int, int] + Spatial resolution of the images. + interpolation : InterpolationMode, optional + Interpolation mode for resizing. (default: InterpolationMode.BILINEAR) + antialias : bool, optional + Whether to use antialiasing. (default: True) + """ + + def __init__( + self, + image_shape: tuple[int, int], + interpolation: InterpolationMode = InterpolationMode.BILINEAR, + antialias: bool = True, + ) -> None: + super().__init__() + self._image_shape = image_shape + self._interpolation = interpolation + self._antialias = antialias + + self._resizer = Resize( + size=self._image_shape, + interpolation=self._interpolation, + antialias=self._antialias + ) + + def send(self, images: torch.Tensor) -> torch.Tensor: + raise RuntimeError( + "FixedResolutionDomain is not supposed to be used for sending images " \ + "because the internal image resolution could not be determined." + ) + + def receive(self, images: torch.Tensor) -> torch.Tensor: + return self._resizer(images) diff --git a/bdpy/dl/torch/torch.py b/bdpy/dl/torch/torch.py index 66825d42..61e0f217 100644 --- a/bdpy/dl/torch/torch.py +++ b/bdpy/dl/torch/torch.py @@ -1,8 +1,11 @@ '''PyTorch module.''' +from __future__ import annotations + from typing import Iterable, List, Dict, Union, Tuple, Any, Callable, Optional from collections import OrderedDict import os +import warnings import numpy as np from PIL import Image @@ -56,10 +59,10 @@ def __init__( layer_object = models._parse_layer_name(self._encoder, layer) layer_object.register_forward_hook(self._extractor) - def __call__(self, x: _tensor_t) -> Dict[str, _tensor_t]: + def __call__(self, x: _tensor_t) -> Dict[str, np.ndarray] | Dict[str, torch.Tensor]: return self.run(x) - def run(self, x: _tensor_t) -> Dict[str, _tensor_t]: + def run(self, x: _tensor_t) -> Dict[str, np.ndarray] | Dict[str, torch.Tensor]: '''Extract feature activations from the specified layers. Parameters @@ -82,17 +85,17 @@ def run(self, x: _tensor_t) -> Dict[str, _tensor_t]: self._encoder.forward(xt) - features: Dict[str, _tensor_t] = { + features: Dict[str, torch.Tensor] = { layer: self._extractor.outputs[i] for i, layer in enumerate(self.__layers) } - if self.__detach: - features = { - k: v.cpu().detach().numpy() - for k, v in features.items() - } + if not self.__detach: + return features - return features + return { + k: v.cpu().detach().numpy() + for k, v in features.items() + } def __del__(self): ''' @@ -182,6 +185,13 @@ def __init__( - Images are converted to RGB. Alpha channels in RGBA images are ignored. ''' + warnings.warn( + "dl.torch.torch.ImageDataset is deprecated. Please consider using " \ + "bdpy.dl.torch.dataset.ImageDataset instead.", + DeprecationWarning, + stacklevel=2 + ) + self.transform = transform # Custom transforms self.__shape = shape diff --git a/bdpy/recon/torch/modules/__init__.py b/bdpy/recon/torch/modules/__init__.py new file mode 100644 index 00000000..7b37bc7e --- /dev/null +++ b/bdpy/recon/torch/modules/__init__.py @@ -0,0 +1,4 @@ +from .encoder import build_encoder, BaseEncoder +from .generator import build_generator, BaseGenerator +from .latent import ArbitraryLatent, BaseLatent +from .critic import TargetNormalizedMSE, BaseCritic diff --git a/bdpy/recon/torch/modules/critic.py b/bdpy/recon/torch/modules/critic.py new file mode 100644 index 00000000..b43d5ba4 --- /dev/null +++ b/bdpy/recon/torch/modules/critic.py @@ -0,0 +1,187 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Dict, Iterable + +import torch +import torch.nn as nn + +from bdpy.task.callback import CallbackHandler, BaseCallback + + +_FeatureType = Dict[str, torch.Tensor] + + +class BaseCritic(ABC): + """Critic network module.""" + + def __init__(self, callbacks: BaseCallback | Iterable[BaseCallback] | None = None) -> None: + self._callback_handler = CallbackHandler(callbacks) + + def __call__(self, features: _FeatureType, target_features: _FeatureType) -> torch.Tensor: + """Call self.compare. + + Parameters + ---------- + features : dict[str, torch.Tensor] + Features indexed by the layer names. + target_features : dict[str, torch.Tensor] + Target features indexed by the layer names. + + Returns + ------- + torch.Tensor + Loss value. + """ + return self.evaluate(features, target_features) + + @abstractmethod + def evaluate( + self, + features: _FeatureType, + target_features: _FeatureType, + ) -> torch.Tensor: + """Compute the total loss value given the features and the target features. + + Parameters + ---------- + features : dict[str, torch.Tensor] + Features indexed by the layer names. + target_features : dict[str, torch.Tensor] + Target features indexed by the layer names. + + Returns + ------- + torch.Tensor + Loss value. + """ + pass + + +class NNModuleCritic(BaseCritic, nn.Module): + """Critic network module uses __call__ method of nn.Module.""" + def __init__(self, callbacks: BaseCallback | Iterable[BaseCallback] | None = None) -> None: + BaseCritic.__init__(self, callbacks) + nn.Module.__init__(self) + + def __call__(self, features: _FeatureType, target_features: _FeatureType) -> torch.Tensor: + return nn.Module.__call__(self, features, target_features) + + def forward(self, features: _FeatureType, target_features: _FeatureType) -> torch.Tensor: + return self.evaluate(features, target_features) + + +class LayerWiseAverageCritic(NNModuleCritic): + """Compute the average of the layer-wise loss values.""" + + def evaluate( + self, + features: _FeatureType, + target_features: _FeatureType, + ) -> torch.Tensor: + """Compute the total loss value given the features and the target features. + + Parameters + ---------- + features : dict[str, torch.Tensor] + Features indexed by the layer names. + target_features : dict[str, torch.Tensor] + Target features indexed by the layer names. + + Returns + ------- + torch.Tensor + Loss value. + """ + loss = 0.0 + counts = 0 + for layer_name, feature in features.items(): + target_feature = target_features[layer_name] + layer_wise_loss = self.compare_layer( + feature, target_feature, layer_name=layer_name + ) + self._callback_handler.fire( + "on_layerwise_loss_calculated", + layer_name=layer_name, + layer_loss=layer_wise_loss, + ) + loss += layer_wise_loss + counts += 1 + return loss / counts + + @abstractmethod + def compare_layer( + self, feature: torch.Tensor, target_feature: torch.Tensor, layer_name: str + ) -> torch.Tensor: + """Loss function per layer. + + Parameters + ---------- + feature : torch.Tensor + Feature tensor of the layer specified by `layer_name`. + target_feature : torch.Tensor + Target feature tensor of the layer specified by `layer_name`. + layer_name : str + Layer name. + + Returns + ------- + torch.Tensor + Loss value of the layer specified by `layer_name`. + """ + pass + + +class MSE(LayerWiseAverageCritic): + """MSE loss.""" + + def compare_layer( + self, feature: torch.Tensor, target_feature: torch.Tensor, layer_name: str + ) -> torch.Tensor: + """Loss function per layer. + + Parameters + ---------- + feature : torch.Tensor + Feature tensor of the layer specified by `layer_name`. + target_feature : torch.Tensor + Target feature tensor of the layer specified by `layer_name`. + layer_name : str + Layer name. + + Returns + ------- + torch.Tensor + Loss value of the layer specified by `layer_name`. + """ + return ((feature - target_feature) ** 2).sum( + dim=tuple(range(1, feature.ndim)) + ) + + +class TargetNormalizedMSE(LayerWiseAverageCritic): + """MSE loss divided by the squared norm of the target feature.""" + + def compare_layer( + self, feature: torch.Tensor, target_feature: torch.Tensor, layer_name: str + ) -> torch.Tensor: + """Loss function per layer. + + Parameters + ---------- + feature : torch.Tensor + Feature tensor of the layer specified by `layer_name`. + target_feature : torch.Tensor + Target feature tensor of the layer specified by `layer_name`. + layer_name : str + Layer name. + + Returns + ------- + torch.Tensor + Loss value of the layer specified by `layer_name`. + """ + squared_norm = (target_feature ** 2).sum(dim=tuple(range(1, target_feature.ndim))) + return ((feature - target_feature) ** 2).sum( + dim=tuple(range(1, feature.ndim)) + ) / squared_norm diff --git a/bdpy/recon/torch/modules/encoder.py b/bdpy/recon/torch/modules/encoder.py new file mode 100644 index 00000000..ec53bb12 --- /dev/null +++ b/bdpy/recon/torch/modules/encoder.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Iterable + +import torch +import torch.nn as nn +from bdpy.dl.torch import FeatureExtractor +from bdpy.dl.torch.domain import Domain, InternalDomain + + +class BaseEncoder(ABC): + """Encoder network module.""" + + @abstractmethod + def encode(self, images: torch.Tensor) -> dict[str, torch.Tensor]: + """Encode images as a hierarchical feature representation. + + Parameters + ---------- + images : torch.Tensor + Images. + + Returns + ------- + dict[str, torch.Tensor] + Features indexed by the layer names. + """ + pass + + def __call__(self, images: torch.Tensor) -> dict[str, torch.Tensor]: + """Call self.encode. + + Parameters + ---------- + images : torch.Tensor + Images on the libraries internal domain. + + Returns + ------- + dict[str, torch.Tensor] + Features indexed by the layer names. + """ + return self.encode(images) + + +class NNModuleEncoder(BaseEncoder, nn.Module): + """Encoder network module subclassed from nn.Module.""" + + def forward(self, images: torch.Tensor) -> dict[str, torch.Tensor]: + """Call self.encode. + + Parameters + ---------- + images : torch.Tensor + Images on the library's internal domain. + + Returns + ------- + dict[str, torch.Tensor] + Features indexed by the layer names. + """ + return self.encode(images) + + def __call__(self, images: torch.Tensor) -> dict[str, torch.Tensor]: + return nn.Module.__call__(self, images) + + +class SimpleEncoder(NNModuleEncoder): + """Encoder network module with a naive feature extractor. + + Parameters + ---------- + feature_network : nn.Module + Feature network. This network should have a method `forward` that takes + an image tensor and propagates it through the network. + layer_names : list[str] + Layer names to extract features from. + domain : Domain, optional + Domain of the input stimuli to receive. (default: InternalDomain()) + + Examples + -------- + >>> import torch + >>> import torch.nn as nn + >>> from bdpy.recon.torch.modules.encoder import SimpleEncoder + >>> feature_network = nn.Sequential( + ... nn.Conv2d(3, 3, 3), + ... nn.ReLU(), + ... ) + >>> encoder = SimpleEncoder(feature_network, ['[0]']) + >>> image = torch.randn(1, 3, 64, 64) + >>> features = encoder(image) + >>> features['[0]'].shape + torch.Size([1, 3, 62, 62]) + """ + + def __init__( + self, + feature_network: nn.Module, + layer_names: Iterable[str], + domain: Domain = InternalDomain(), + ) -> None: + super().__init__() + self._feature_extractor = FeatureExtractor( + encoder=feature_network, layers=layer_names, detach=False, device=None + ) + self._domain = domain + self._feature_network = self._feature_extractor._encoder + + def encode(self, images: torch.Tensor) -> dict[str, torch.Tensor]: + """Encode images as a hierarchical feature representation. + + Parameters + ---------- + images : torch.Tensor + Images on the libraries internal domain. + + Returns + ------- + dict[str, torch.Tensor] + Features indexed by the layer names. + """ + images = self._domain.receive(images) + return self._feature_extractor(images) + + +def build_encoder( + feature_network: nn.Module, + layer_names: Iterable[str], + domain: Domain = InternalDomain(), +) -> BaseEncoder: + """Build an encoder network with a naive feature extractor. + + The function builds an encoder module from a feature network that takes + images on its own domain as input and processes them. The encoder module + receives images on the library's internal domain and returns features on the + library's internal domain indexed by `layer_names`. `domain` is used to + convert the input images to the feature network's domain from the library's + internal domain. + + Parameters + ---------- + feature_network : nn.Module + Feature network. This network should have a method `forward` that takes + an image tensor and propagates it through the network. The images should + be on the network's own domain. + layer_names : list[str] + Layer names to extract features from. + domain : Domain, optional + Domain of the input stimuli to receive (default: InternalDomain()). + One needs to specify the domain that corresponds to the feature network's + input domain. + + Returns + ------- + BaseEncoder + Encoder network. + + Examples + -------- + >>> import torch + >>> import torch.nn as nn + >>> from bdpy.recon.torch.modules.encoder import build_encoder + >>> feature_network = nn.Sequential( + ... nn.Conv2d(3, 3, 3), + ... nn.ReLU(), + ... ) + >>> encoder = build_encoder(feature_network, layer_names=['[0]']) + >>> image = torch.randn(1, 3, 64, 64) + >>> features = encoder(image) + >>> features['[0]'].shape + torch.Size([1, 3, 62, 62]) + """ + return SimpleEncoder(feature_network, layer_names, domain) diff --git a/bdpy/recon/torch/modules/generator.py b/bdpy/recon/torch/modules/generator.py new file mode 100644 index 00000000..98f92109 --- /dev/null +++ b/bdpy/recon/torch/modules/generator.py @@ -0,0 +1,309 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Callable, Iterator +import warnings + +import torch +import torch.nn as nn +from torch.nn.parameter import Parameter +from bdpy.dl.torch.domain import Domain, InternalDomain + + +def _get_reset_module_fn(module: nn.Module) -> Callable[[], None] | None: + """Get the function to reset the parameters of the module.""" + reset_parameters = getattr(module, "reset_parameters", None) + if callable(reset_parameters): + return reset_parameters + # NOTE: This is needed for nn.MultiheadAttention + reset_parameters = getattr(module, "_reset_parameters", None) + if callable(reset_parameters): + return reset_parameters + return None + + +@torch.no_grad() +def call_reset_parameters(module: nn.Module) -> None: + """Reset the parameters of the module.""" + warnings.warn( + "`call_reset_parameters` calls the instance method named `reset_parameters` " \ + "or `_reset_parameters` of the module. This method does not guarantee that " \ + "all the parameters of the module are reset. Please use this method with " \ + "caution.", + UserWarning, + stacklevel=2, + ) + reset_parameters = _get_reset_module_fn(module) + if reset_parameters is not None: + reset_parameters() + + +class BaseGenerator(ABC): + """Generator module.""" + + @abstractmethod + def reset_states(self) -> None: + """Reset the state of the generator.""" + pass + + @abstractmethod + def parameters(self, recurse: bool = True) -> Iterator[nn.Parameter]: + """Return the parameters of the generator.""" + pass + + @abstractmethod + def generate(self, latent: torch.Tensor) -> torch.Tensor: + """Generate image given latent variable. + + Parameters + ---------- + latent : torch.Tensor + Latent variable. + + Returns + ------- + torch.Tensor + Generated image on the libraries internal domain. + """ + pass + + def __call__(self, latent: torch.Tensor) -> torch.Tensor: + """Call self.generate. + + Parameters + ---------- + latent : torch.Tensor + Latent vector. + + Returns + ------- + torch.Tensor + Generated image. The generated images must be in the range [0, 1]. + """ + return self.generate(latent) + + +class NNModuleGenerator(BaseGenerator, nn.Module): + """Generator module uses __call__ method and parameters method of nn.Module.""" + + def parameters(self, recurse: bool = True) -> Iterator[nn.Parameter]: + return nn.Module.parameters(self, recurse=recurse) + + def forward(self, latent: torch.Tensor) -> torch.Tensor: + return self.generate(latent) + + def __call__(self, latent: torch.Tensor) -> torch.Tensor: + return nn.Module.__call__(self, latent) + + +class BareGenerator(NNModuleGenerator): + """Bare generator module. + + This module does not have any trainable parameters. + + Parameters + ---------- + activation : Callable[[torch.Tensor], torch.Tensor], optional + Activation function to apply to the output of the generator, by default nn.Identity() + + Examples + -------- + >>> import torch + >>> from bdpy.recon.torch.modules.generator import BareGenerator + >>> generator = BareGenerator(activation=torch.sigmoid) + >>> latent = torch.randn(1, 3, 64, 64) + >>> generated_image = generator(latent) + >>> generated_image.shape + torch.Size([1, 3, 64, 64]) + """ + + def __init__(self, activation: Callable[[torch.Tensor], torch.Tensor] = nn.Identity()) -> None: + """Initialize the generator.""" + super().__init__() + self._activation = activation + self._domain = InternalDomain() + + def reset_states(self) -> None: + """Reset the state of the generator.""" + pass + + def generate(self, latent: torch.Tensor) -> torch.Tensor: + """Naively pass the latent vector to the activation function. + + Parameters + ---------- + latent : torch.Tensor + Latent vector. + + Returns + ------- + torch.Tensor + Generated image on the libraries internal domain. + """ + return self._domain.send(self._activation(latent)) + + +class DNNGenerator(NNModuleGenerator): + """DNN generator module. + + This module has the generator network as a submodule and its parameters are + trainable. + + Parameters + ---------- + generator_network : nn.Module + Generator network. This network should have a method `forward` that takes + a latent vector and propagates it through the network. + domain : Domain, optional + Domain of the input stimuli to receive. (default: InternalDomain()) + reset_fn : Callable[[nn.Module], None], optional + Function to reset the parameters of the generator network, by default + call_reset_parameters. + + Examples + -------- + >>> import torch + >>> import torch.nn as nn + >>> from bdpy.recon.torch.modules.generator import DNNGenerator + >>> generator_network = nn.Sequential( + ... nn.ConvTranspose2d(3, 3, 3), + ... nn.ReLU(), + ... ) + >>> generator = DNNGenerator(generator_network) + >>> latent = torch.randn(1, 3, 64, 64) + >>> generated_image = generator(latent) + >>> generated_image.shape + torch.Size([1, 3, 66, 66]) + """ + + def __init__( + self, + generator_network: nn.Module, + domain: Domain = InternalDomain(), + reset_fn: Callable[[nn.Module], None] = call_reset_parameters, + ) -> None: + """Initialize the generator.""" + super().__init__() + self._generator_network = generator_network + self._domain = domain + self._reset_fn = reset_fn + + def reset_states(self) -> None: + """Reset the state of the generator.""" + self._generator_network.apply(self._reset_fn) + + def generate(self, latent: torch.Tensor) -> torch.Tensor: + """Generate image using the generator network. + + Parameters + ---------- + latent : torch.Tensor + Latent vector. + + Returns + ------- + torch.Tensor + Generated image on the libraries internal domain. + """ + return self._domain.send(self._generator_network(latent)) + + +class FrozenGenerator(DNNGenerator): + """Frozen generator module. + + This module has the generator network as a submodule and its parameters are + frozen. + + Parameters + ---------- + generator_network : nn.Module + Generator network. This network should have a method `forward` that takes + a latent vector and propagates it through the network. + domain : Domain, optional + Domain of the input stimuli to receive. (default: InternalDomain()) + + Examples + -------- + >>> import torch + >>> import torch.nn as nn + >>> from bdpy.recon.torch.modules.generator import FrozenGenerator + >>> generator_network = nn.Sequential( + ... nn.ConvTranspose2d(3, 3, 3), + ... nn.ReLU(), + ... ) + >>> generator = FrozenGenerator(generator_network) + >>> latent = torch.randn(1, 3, 64, 64) + >>> generated_image = generator(latent) + >>> generated_image.shape + torch.Size([1, 3, 66, 66]) + """ + + def __init__( + self, + generator_network: nn.Module, + domain: Domain = InternalDomain(), + ) -> None: + """Initialize the generator.""" + super().__init__(generator_network, domain=domain, reset_fn=lambda _: None) + self._generator_network.eval() + + def reset_states(self) -> None: + """Reset the state of the generator.""" + pass + + def parameters(self, recurse: bool = True) -> Iterator[Parameter]: + """Return an empty iterator.""" + return iter([]) + + +def build_generator( + generator_network: nn.Module, + domain: Domain = InternalDomain(), + reset_fn: Callable[[nn.Module], None] = call_reset_parameters, + frozen: bool = True, +) -> BaseGenerator: + """Build a generator module. + + This function builds a generator module from a generator network that takes + a latent vector as an input and returns an image on its own domain. One + needs to specify the domain of the generator network. + + Parameters + ---------- + generator_network : nn.Module + Generator network. This network should have a method `forward` that takes + a latent vector and returns an image on its own domain. + domain : Domain, optional + Domain of the input images to receive. (default: InternalDomain()). + One needs to specify the domain that corresponds to the generator + network's output domain. + reset_fn : Callable[[nn.Module], None], optional + Function to reset the parameters of the generator network, by default + call_reset_parameters. + frozen : bool, optional + Whether to freeze the parameters of the generator network, by default True. + + Returns + ------- + BaseGenerator + Generator module. + + Examples + -------- + >>> import torch + >>> import torch.nn as nn + >>> from bdpy.recon.torch.modules.generator import build_generator + >>> generator_network = nn.Sequential( + ... nn.ConvTranspose2d(3, 3, 3), + ... nn.ReLU(), + ... ) + >>> generator = build_generator(generator_network) + >>> latent = torch.randn(1, 3, 64, 64) + >>> generated_image = generator(latent) + >>> generated_image.shape + torch.Size([1, 3, 66, 66]) + """ + if frozen: + return FrozenGenerator(generator_network, domain=domain) + else: + return DNNGenerator(generator_network, domain=domain, reset_fn=reset_fn) diff --git a/bdpy/recon/torch/modules/latent.py b/bdpy/recon/torch/modules/latent.py new file mode 100644 index 00000000..28195c70 --- /dev/null +++ b/bdpy/recon/torch/modules/latent.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Callable, Iterator + +import torch +import torch.nn as nn + + +class BaseLatent(ABC): + """Latent variable module.""" + + @abstractmethod + def reset_states(self) -> None: + """Reset the state of the latent variable.""" + pass + + @abstractmethod + def parameters(self, recurse: bool = True) -> Iterator[nn.Parameter]: + """Return the parameters of the latent variable.""" + pass + + @abstractmethod + def generate(self) -> torch.Tensor: + """Generate a latent variable. + + Returns + ------- + torch.Tensor + Latent variable. + """ + pass + + def __call__(self) -> torch.Tensor: + """Call self.generate. + + Returns + ------- + torch.Tensor + Latent variable. + """ + return self.generate() + + +class NNModuleLatent(BaseLatent, nn.Module): + """Latent variable module uses __call__ method and parameters method of nn.Module.""" + + def parameters(self, recurse: bool = True) -> Iterator[nn.Parameter]: + return nn.Module.parameters(self, recurse=recurse) + + def __call__(self) -> torch.Tensor: + return nn.Module.__call__(self) + + def forward(self) -> torch.Tensor: + return self.generate() + + +class ArbitraryLatent(NNModuleLatent): + """Latent variable with arbitrary shape and initialization function. + + Parameters + ---------- + shape : tuple[int, ...] + Shape of the latent variable including the batch dimension. + init_fn : Callable[[torch.Tensor], None] + Function to initialize the latent variable. + + Examples + -------- + >>> from functools import partial + >>> import torch + >>> import torch.nn as nn + >>> from bdpy.recon.torch.modules.latent import ArbitraryLatent + >>> latent = ArbitraryLatent((1, 3, 64, 64), partial(nn.init.normal_, mean=0, std=1)) + >>> latent().shape + torch.Size([1, 3, 64, 64]) + """ + + def __init__(self, shape: tuple[int, ...], init_fn: Callable[[torch.Tensor], None]) -> None: + super().__init__() + self._shape = shape + self._init_fn = init_fn + self._latent = nn.Parameter(torch.empty(shape)) + + def reset_states(self) -> None: + """Reset the state of the latent variable.""" + self._init_fn(self._latent) + + def generate(self) -> torch.Tensor: + """Generate a latent variable. + + Returns + ------- + torch.Tensor + Latent variable. + """ + return self._latent diff --git a/bdpy/recon/torch/task/__init__.py b/bdpy/recon/torch/task/__init__.py new file mode 100644 index 00000000..a5292096 --- /dev/null +++ b/bdpy/recon/torch/task/__init__.py @@ -0,0 +1 @@ +from .inversion import FeatureInversionTask \ No newline at end of file diff --git a/bdpy/recon/torch/task/inversion.py b/bdpy/recon/torch/task/inversion.py new file mode 100644 index 00000000..8a3b40b8 --- /dev/null +++ b/bdpy/recon/torch/task/inversion.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +from typing import Dict, Iterable, Callable + +from itertools import chain + +import torch + +from ..modules import BaseEncoder, BaseGenerator, BaseLatent, BaseCritic +from bdpy.task import BaseTask +from bdpy.task.callback import BaseCallback, unused, _validate_callback + +FeatureType = Dict[str, torch.Tensor] + + +def _apply_to_features( + fn: Callable[[torch.Tensor], torch.Tensor], features: FeatureType +) -> FeatureType: + return {k: fn(v) for k, v in features.items()} + + +class FeatureInversionCallback(BaseCallback): + """Callback for feature inversion task. + + As a design principle, the callback functions must not have any side effects + on the task results. It should be used only for logging, visualization, + etc. Please refer to `bdpy.util.callback.BaseCallback` for details of the + usage of callbacks. + """ + + def __init__(self) -> None: + super().__init__(base_class=FeatureInversionCallback) + + @unused + def on_task_start(self) -> None: + """Callback on task start.""" + pass + + @unused + def on_iteration_start(self, *, step: int) -> None: + """Callback on iteration start.""" + pass + + @unused + def on_image_generated(self, *, step: int, image: torch.Tensor) -> None: + """Callback on image generated.""" + pass + + @unused + def on_layerwise_loss_calculated( + self, *, layer_loss: torch.Tensor, layer_name: str + ) -> None: + """Callback on layerwise loss calculated.""" + pass + + @unused + def on_loss_calculated(self, *, step: int, loss: torch.Tensor) -> None: + """Callback on loss calculated.""" + pass + + @unused + def on_iteration_end(self, *, step: int) -> None: + """Called at the end of each iteration.""" + pass + + @unused + def on_task_end(self) -> None: + """Callback on task end.""" + pass + + +class CUILoggingCallback(FeatureInversionCallback): + """Callback for logging on CUI. + + Parameters + ---------- + interval : int, optional + Logging interval, by default 1. If `interval` is 1, the callback logs + every iteration. + total_steps : int, optional + Total number of iterations, by default -1. If `total_steps` is -1, + the callback does not show the total number of iterations. + """ + + def __init__(self, interval: int = 1, total_steps: int = -1) -> None: + super().__init__() + self._interval = interval + self._total_steps = total_steps + self._loss: int | float = -1 + + def _step_str(self, step: int) -> str: + if self._total_steps > 0: + return f"{step+1}/{self._total_steps}" + else: + return f"{step+1}" + + def on_loss_calculated(self, *, step: int, loss: torch.Tensor) -> None: + self._loss = loss.item() + + def on_iteration_end(self, *, step: int) -> None: + if step % self._interval == 0: + print(f"Step: [{self._step_str(step)}], Loss: {self._loss:.4f}") + + +class FeatureInversionTask(BaseTask): + """Feature inversion Task. + + Parameters + ---------- + encoder : BaseEncoder + Encoder module. + generator : BaseGenerator + Generator module. + latent : BaseLatent + Latent variable module. + critic : BaseCritic + Critic module. + optimizer : torch.optim.Optimizer + Optimizer. + scheduler : torch.optim.lr_scheduler.LRScheduler, optional + Learning rate scheduler, by default None. + num_iterations : int, optional + Number of iterations, by default 1. + callbacks : FeatureInversionCallback | Iterable[FeatureInversionCallback] | None, optional + Callbacks, by default None. Please refer to `bdpy.util.callback.BaseCallback` + and `bdpy.recon.torch.task.FeatureInversionCallback` for details. + + Examples + -------- + >>> import torch + >>> import torch.nn as nn + >>> from bdpy.recon.torch.task import FeatureInversionTask + >>> from bdpy.recon.torch.modules import build_encoder, build_generator, ArbitraryLatent, TargetNormalizedMSE + >>> encoder = build_encoder(...) + >>> generator = build_generator(...) + >>> latent = ArbitraryLatent(...) + >>> critic = TargetNormalizedMSE(...) + >>> optimizer = torch.optim.Adam(latent.parameters()) + >>> task = FeatureInversionTask( + ... encoder, generator, latent, critic, optimizer, num_iterations=200, + ... ) + >>> target_features = encoder(target_image) + >>> reconstructed_image = task(target_features) + """ + + def __init__( + self, + encoder: BaseEncoder, + generator: BaseGenerator, + latent: BaseLatent, + critic: BaseCritic, + optimizer: torch.optim.Optimizer, + scheduler: torch.optim.lr_scheduler.LRScheduler = None, + num_iterations: int = 1, + callbacks: FeatureInversionCallback + | Iterable[FeatureInversionCallback] + | None = None, + ) -> None: + super().__init__(callbacks) + self._encoder = encoder + self._generator = generator + self._latent = latent + self._critic = critic + self._optimizer = optimizer + self._scheduler = scheduler + + self._num_iterations = num_iterations + + def run( + self, + target_features: FeatureType, + ) -> torch.Tensor: + """Run feature inversion given target features. + + Parameters + ---------- + target_features : FeatureType + Target features. + + Returns + ------- + torch.Tensor + Reconstructed images on the libraries internal domain. + """ + self._callback_handler.fire("on_task_start") + self.reset_states() + for step in range(self._num_iterations): + self._callback_handler.fire("on_iteration_start", step=step) + self._optimizer.zero_grad() + + latent = self._latent() + generated_image = self._generator(latent) + self._callback_handler.fire( + "on_image_generated", step=step, image=generated_image.clone().detach() + ) + + features = self._encoder(generated_image) + + loss = self._critic(features, target_features) + self._callback_handler.fire( + "on_loss_calculated", step=step, loss=loss.clone().detach() + ) + loss.backward() + + self._optimizer.step() + if self._scheduler is not None: + self._scheduler.step() + + self._callback_handler.fire("on_iteration_end", step=step) + + generated_image = self._generator(self._latent()).detach() + + self._callback_handler.fire("on_task_end") + return generated_image + + def reset_states(self) -> None: + """Reset the state of the task.""" + self._generator.reset_states() + self._latent.reset_states() + self._optimizer = self._optimizer.__class__( + chain( + self._generator.parameters(), + self._latent.parameters(), + ), + **self._optimizer.defaults, + ) diff --git a/bdpy/task/__init__.py b/bdpy/task/__init__.py new file mode 100644 index 00000000..b1ee439e --- /dev/null +++ b/bdpy/task/__init__.py @@ -0,0 +1 @@ +from .core import BaseTask diff --git a/bdpy/task/callback.py b/bdpy/task/callback.py new file mode 100644 index 00000000..93f097f4 --- /dev/null +++ b/bdpy/task/callback.py @@ -0,0 +1,265 @@ +from __future__ import annotations + +from typing import Callable, Type, Any, Iterable, TypeVar, Generic +from typing_extensions import Annotated, ParamSpec + +from collections import defaultdict +from functools import wraps + + +_P = ParamSpec("_P") +_Unused = Annotated[None, "unused"] + + +def _is_unused(fn: Callable) -> bool: + if not hasattr(fn, "__annotations__"): + return False + return_type: Type | None = fn.__annotations__.get("return", None) + if return_type is None: + return False + return return_type == _Unused + + +def unused(fn: Callable[_P, Any]) -> Callable[_P, _Unused]: + """Decorate a function to raise an error when called. + + This decorator marks a function as unused and raises an error when called. + The type of the return value is changed to `Annotated[None, "unused"]`. + + Parameters + ---------- + fn : Callable + Function to decorate. + + Returns + ------- + Callable + Decorated function. + + Examples + -------- + >>> @unused + ... def f(a: int, b: int, c: int = 0) -> int: + ... return a + b + c + ... + >>> f(1, 2, 3) + Traceback (most recent call last): + ... + RuntimeError: Function is decorated with @unused and must not be called. + """ + + @wraps(fn) # NOTE: preserve name, docstring, etc. of the original function + def _unused(*args: _P.args, **kwargs: _P.kwargs) -> _Unused: + raise RuntimeError( + f"Function {fn} is decorated with @unused and must not be called." + ) + + # NOTE: change the return type to Unused + _unused.__annotations__["return"] = _Unused + + return _unused + + +def _validate_callback(callback: BaseCallback, base_class: Type[BaseCallback]) -> None: + """Validate a callback. + + Parameters + ---------- + callback : BaseCallback + Callback to validate. + base_class : Type[BaseCallback] + Base class of the callback. + + Raises + ------ + TypeError + If the callback is not an instance of the base class. + ValueError + If the callback has an event type that is not acceptable. + + Examples + -------- + >>> class TaskBaseCallback(BaseCallback): + ... @unused + ... def on_task_start(self): + ... pass + ... + ... @unused + ... def on_task_end(self): + ... pass + ... + >>> class SomeTaskCallback(TaskBaseCallback): + ... def on_unacceptable_event(self): + ... # do something + ... + >>> callback = SomeTaskCallback() + >>> _validate_callback(callback, TaskBaseCallback) + Traceback (most recent call last): + ... + ValueError: on_unacceptable_event is not an acceptable event type. ... + """ + + if not isinstance(callback, base_class): + raise TypeError( + f"Callback must be an instance of {base_class}, not {type(callback)}." + ) + acceptable_events = [] + for event_type in dir(base_class): + if event_type.startswith("on_") and callable(getattr(base_class, event_type)): + acceptable_events.append(event_type) + for event_type in dir(callback): + if not ( + event_type.startswith("on_") and callable(getattr(callback, event_type)) + ): + continue + if event_type not in acceptable_events: + raise ValueError( + f"{event_type} is not an acceptable event type. " + f"Acceptable event types are {acceptable_events}. " + f"Please refer to the documentation of {base_class.__name__} for the list of acceptable event types." + ) + + +class BaseCallback: + """Base class for callbacks. + + Callbacks are used to hook into the task and execute custom functions + at specific events. Callback functions must be defined as methods of the + callback classes. The callback functions must be named as "on_". + As a design principle, the callback functions must not have any side effects + on the task results. It should be used only for logging, visualization, etc. + + For example, the following callback class logs the start and end of the task. + + >>> class Callback(BaseCallback): + ... def on_task_start(self): + ... print("Task started.") + ... + ... def on_task_end(self): + ... print("Task ended.") + ... + >>> callback = Callback() + >>> some_task = SomeTask() # Initialize a task object + >>> some_task.register_callback(callback) + >>> outputs = some_task(inputs) # Run the task + Task started. + Task ended. + + The set of available events that can be hooked into depends on the task. + See the base class of the corresponding task for the list of all events. + `@unused` decorator can be used to mark a callback function as unused, so + that the callback handler does not fire the function. + """ + def __init__(self, base_class: Type[BaseCallback] | None = None) -> None: + if base_class is None: + base_class = BaseCallback + _validate_callback(self, base_class) + + @unused + def on_task_start(self) -> None: + """Callback on task start.""" + pass + + @unused + def on_task_end(self) -> None: + """Callback on task end.""" + pass + + +_CallbackType = TypeVar("_CallbackType", bound=BaseCallback) + + +class CallbackHandler(Generic[_CallbackType]): + """Callback handler. + + This class manages the callback objects and fires the callback functions + registered to the event type. The callback functions must be defined as + methods of the callback classes. The callback functions must be named as + "on_". + + Parameters + ---------- + callbacks : BaseCallback | Iterable[BaseCallback] | None, optional + Callbacks to register, by default None + + Examples + -------- + >>> class Callback(BaseCallback): + ... def __init__(self, name): + ... self._name = name + ... + ... def on_task_start(self): + ... print(f"Task started (name={self._name}).") + ... + ... def on_task_end(self): + ... print(f"Task ended (name={self._name}).") + ... + >>> handler = CallbackHandler([Callback("A"), Callback("B")]) + >>> handler.fire("on_task_start") + Task started (name=A). + Task started (name=B). + >>> handler.fire("on_task_end") + Task ended (name=A). + Task ended (name=B). + """ + + _callbacks: list[_CallbackType] + _registered_functions: defaultdict[str, list[Callable]] + + def __init__( + self, callbacks: _CallbackType | Iterable[_CallbackType] | None = None + ) -> None: + self._callbacks = [] + self._registered_functions = defaultdict(list) + if callbacks is not None: + if not isinstance(callbacks, Iterable): + callbacks = [callbacks] + for callback in callbacks: + self.register(callback) + + def register(self, callback: _CallbackType) -> None: + """Register a callback. + + Parameters + ---------- + callback : BaseCallback + Callback to register. + + Raises + ------ + TypeError + If the callback is not an instance of BaseCallback. + """ + if not isinstance(callback, BaseCallback): + raise TypeError( + f"Callback must be an instance of BaseCallback, not {type(callback)}." + ) + + self._callbacks.append(callback) + for event_type in dir(callback): + callback_method = getattr(callback, event_type) + if not callable(callback_method): + continue + if not event_type.startswith("on_"): + continue + if _is_unused(callback_method): + continue + self._registered_functions[event_type].append(callback_method) + + def fire(self, event_type: str, **kwargs: Any) -> None: + """Execute the callback functions registered to the event type. + + Parameters + ---------- + event_type : str + Event type to fire, which must start with "on_". + kwargs : dict[str, Any] + Keyword arguments to pass to the callback functions. + + Raises + ------ + KeyError + If the event type is not registered. + """ + for callback_method in self._registered_functions[event_type]: + callback_method(**kwargs) diff --git a/bdpy/task/core.py b/bdpy/task/core.py new file mode 100644 index 00000000..d345a43e --- /dev/null +++ b/bdpy/task/core.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod + +from typing import Iterable, Any, TypeVar, Generic + +from bdpy.task.callback import CallbackHandler, BaseCallback + + +_CallbackType = TypeVar("_CallbackType", bound=BaseCallback) + + +class BaseTask(ABC, Generic[_CallbackType]): + """Base class for tasks. + + Parameters + ---------- + callbacks : BaseCallback | Iterable[BaseCallback] | None + Callbacks to register. If `None`, no callbacks are registered. + + Attributes + ---------- + _callback_handler : CallbackHandler + Callback handler. + + Notes + ----- + This class is designed to be used as a base class for tasks. The task + implementation should override the `__call__` method. The actual interface + of `__call__` depends on the task. For example, the task may take a single + input and return a single output, or it may take multiple inputs and return + multiple outputs. The task may also take keyword arguments. Please refer to + the documentation of the specific task for details. + """ + + _callback_handler: CallbackHandler[_CallbackType] + + def __init__( + self, callbacks: _CallbackType | Iterable[_CallbackType] | None = None + ) -> None: + self._callback_handler = CallbackHandler(callbacks) + + def __call__(self, *inputs, **parameters) -> Any: + """Run the task.""" + return self.run(*inputs, **parameters) + + @abstractmethod + def run(self, *inputs, **parameters) -> Any: + """Run the task.""" + pass + + def register_callback(self, callback: _CallbackType) -> None: + """Register a callback. + + Parameters + ---------- + callback : BaseCallback + Callback to register. + """ + self._callback_handler.register(callback) diff --git a/pyproject.toml b/pyproject.toml index 142e9530..7e959a49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ caffe = [ ] torch = [ "torch", + "torchvision", "Pillow" ] fig = [ diff --git a/tests/dl/torch/domain/__init__.py b/tests/dl/torch/domain/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/dl/torch/domain/test_core.py b/tests/dl/torch/domain/test_core.py new file mode 100644 index 00000000..092c3487 --- /dev/null +++ b/tests/dl/torch/domain/test_core.py @@ -0,0 +1,138 @@ +"""Tests for bdpy.dl.torch.domain.core.""" + +import unittest +from bdpy.dl.torch.domain import core as core_module + + +class DummyAddDomain(core_module.Domain): + def send(self, num): + return num + 1 + + def receive(self, num): + return num - 1 + + +class DummyDoubleDomain(core_module.Domain): + def send(self, num): + return num * 2 + + def receive(self, num): + return num // 2 + + +class DummyUpperCaseDomain(core_module.Domain): + def send(self, text): + return text.upper() + + def receive(self, value): + return value.lower() + + +class TestDomain(unittest.TestCase): + """Tests for bdpy.dl.torch.domain.core.Domain.""" + def setUp(self): + self.domain = DummyAddDomain() + self.original_space_num = 0 + self.internal_space_num = 1 + + def test_instantiation(self): + """Test instantiation.""" + self.assertRaises(TypeError, core_module.Domain) + + def test_send(self): + """test send""" + self.assertEqual(self.domain.send(self.original_space_num), self.internal_space_num) + + def test_receive(self): + """test receive""" + self.assertEqual(self.domain.receive(self.internal_space_num), self.original_space_num) + + def test_invertibility(self): + input_candidates = [-1, 0, 1, 0.5] + for x in input_candidates: + assert x == self.domain.send(self.domain.receive(x)) + assert x == self.domain.receive(self.domain.send(x)) + + +class TestInternalDomain(unittest.TestCase): + """Tests for bdpy.dl.torch.domain.core.InternalDomain.""" + def setUp(self): + self.domain = core_module.InternalDomain() + self.num = 1 + + def test_send(self): + """test send""" + self.assertEqual(self.domain.send(self.num), self.num) + + def test_receive(self): + """test receive""" + self.assertEqual(self.domain.receive(self.num), self.num) + + def test_invertibility(self): + input_candidates = [-1, 0, 1, 0.5] + for x in input_candidates: + assert x == self.domain.send(self.domain.receive(x)) + assert x == self.domain.receive(self.domain.send(x)) + + +class TestIrreversibleDomain(unittest.TestCase): + """Tests for bdpy.dl.torch.domain.core.IrreversibleDomain.""" + def setUp(self): + self.domain = core_module.IrreversibleDomain() + self.num = 1 + + def test_send(self): + """test send""" + self.assertEqual(self.domain.send(self.num), self.num) + + def test_receive(self): + """test receive""" + self.assertEqual(self.domain.receive(self.num), self.num) + + +class TestComposedDomain(unittest.TestCase): + """Tests for bdpy.dl.torch.domain.core.ComposedDomain.""" + def setUp(self): + self.composed_domain = core_module.ComposedDomain([ + DummyDoubleDomain(), + DummyAddDomain(), + ]) + self.original_space_num = 0 + self.internal_space_num = 2 + + def test_send(self): + """test send""" + self.assertEqual(self.composed_domain.send(self.original_space_num), self.internal_space_num) + + def test_receive(self): + """test receive""" + self.assertEqual(self.composed_domain.receive(self.internal_space_num), self.original_space_num) + + +class TestKeyValueDomain(unittest.TestCase): + """Tests for bdpy.dl.torch.domain.core.KeyValueDomain.""" + def setUp(self): + self.key_value_domain = core_module.KeyValueDomain({ + "name": DummyUpperCaseDomain(), + "age": DummyDoubleDomain() + }) + self.original_space_data = {"name": "alice", "age": 30} + self.internal_space_data = {"name": "ALICE", "age": 60} + + def test_send(self): + """test send""" + self.assertEqual(self.key_value_domain.send(self.original_space_data), self.internal_space_data) + + def test_receive(self): + """test receive""" + self.assertEqual(self.key_value_domain.receive(self.internal_space_data), self.original_space_data) + + +if __name__ == "__main__": + #unittest.main() + composed_domain = core_module.ComposedDomain([ + DummyDoubleDomain(), + DummyAddDomain(), + ]) + print(composed_domain.receive(-1)) + print(composed_domain.send(-2)) diff --git a/tests/dl/torch/domain/test_feature_domain.py b/tests/dl/torch/domain/test_feature_domain.py new file mode 100644 index 00000000..970950b3 --- /dev/null +++ b/tests/dl/torch/domain/test_feature_domain.py @@ -0,0 +1,86 @@ +"""Tests for bdpy.dl.torch.domain.feature_domain.""" + +import unittest +import torch +from bdpy.dl.torch.domain import feature_domain as feature_domain_module + + +class TestMethods(unittest.TestCase): + def setUp(self): + self.lnd_tensor = torch.empty((12, 196, 768)) + self.nld_tensor = torch.empty((196, 12, 768)) + + def test_lnd2nld(self): + """test _lnd2nld""" + self.assertEqual(feature_domain_module._lnd2nld(self.lnd_tensor).shape, self.nld_tensor.shape) + + def test_nld2lnd(self): + """test _nld2lnd""" + self.assertEqual(feature_domain_module._nld2lnd(self.nld_tensor).shape, self.lnd_tensor.shape) + + +class TestArbitraryFeatureKeyDomain(unittest.TestCase): + """Tests for bdpy.dl.torch.domain.feature_domain.ArbitraryFeatureKeyDomain.""" + def setUp(self): + self.to_internal_mapping = { + "self_key1": "internal_key1", + "self_key2": "internal_key2" + } + self.to_self_mapping = { + "internal_key1": "self_key1", + "internal_key2": "self_key2" + } + self.features = { + "self_key1": 123, + "self_key2": 456 + } + self.internal_features = { + "internal_key1": 123, + "internal_key2": 456 + } + + def test_send(self): + """test send""" + # when both are specified + domain = feature_domain_module.ArbitraryFeatureKeyDomain( + to_internal=self.to_internal_mapping, + to_self=self.to_self_mapping + ) + self.assertEqual(domain.send(self.features), self.internal_features) + + # when only to_self is specified + domain = feature_domain_module.ArbitraryFeatureKeyDomain( + to_self=self.to_self_mapping + ) + self.assertEqual(domain.send(self.features), self.internal_features) + + # when only to_internal is specified + domain = feature_domain_module.ArbitraryFeatureKeyDomain( + to_internal=self.to_internal_mapping + ) + self.assertEqual(domain.send(self.features), self.internal_features) + + def test_receive(self): + """test receive""" + # when both are specified + domain = feature_domain_module.ArbitraryFeatureKeyDomain( + to_internal=self.to_internal_mapping, + to_self=self.to_self_mapping + ) + self.assertEqual(domain.receive(self.internal_features), self.features) + + # when only to_self is specified + domain = feature_domain_module.ArbitraryFeatureKeyDomain( + to_self=self.to_self_mapping + ) + self.assertEqual(domain.receive(self.internal_features), self.features) + + # when only to_internal is specified + domain = feature_domain_module.ArbitraryFeatureKeyDomain( + to_internal=self.to_internal_mapping + ) + self.assertEqual(domain.receive(self.internal_features), self.features) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/tests/dl/torch/domain/test_image_domain.py b/tests/dl/torch/domain/test_image_domain.py new file mode 100644 index 00000000..6140ab26 --- /dev/null +++ b/tests/dl/torch/domain/test_image_domain.py @@ -0,0 +1,143 @@ +"""Tests for bdpy.dl.torch.domain.image_domain.""" + +import unittest +import torch +import numpy as np +import warnings +from bdpy.dl.torch.domain import image_domain as iamge_domain_module + + +class TestAffineDomain(unittest.TestCase): + """Tests for bdpy.dl.torch.domain.image_domain.AffineDomain""" + def setUp(self): + self.center0d = 0.0 + self.center1d = np.random.randn(3) + self.center2d = np.random.randn(32, 32) + self.center3d = np.random.randn(3, 32, 32) + self.scale0d = 1 + self.scale1d = np.random.randn(3) + self.scale2d = np.random.randn(32, 32) + self.scale3d = np.random.randn(3, 32, 32) + self.image = torch.rand((1, 3, 32, 32)) + + def test_instantiation(self): + """Test instantiation.""" + # Succeeds when center and scale are 0-dimensional + affine_domain = iamge_domain_module.AffineDomain(self.center0d, self.scale0d) + self.assertIsInstance(affine_domain, iamge_domain_module.AffineDomain) + + # Succeeds when center and scale are 1-dimensional + affine_domain = iamge_domain_module.AffineDomain(self.center1d, self.scale1d) + self.assertIsInstance(affine_domain, iamge_domain_module.AffineDomain) + + # Succeeds when center and scale are 3-dimensional + affine_domain = iamge_domain_module.AffineDomain(self.center3d, self.scale3d) + self.assertIsInstance(affine_domain, iamge_domain_module.AffineDomain) + + # Fails when the center is neither 1-dimensional nor 3-dimensional + with self.assertRaises(ValueError): + iamge_domain_module.AffineDomain(self.center2d, self.scale0d) + + # Fails when the scale is neither 1-dimensional nor 3-dimensional + with self.assertRaises(ValueError): + iamge_domain_module.AffineDomain(self.center0d, self.scale2d) + + def test_send_and_receive(self): + """Test send and receive""" + # when 0d + affine_domain = iamge_domain_module.AffineDomain(self.center0d, self.scale0d) + transformed_image = affine_domain.send(self.image) + center0d = torch.from_numpy(np.array([self.center0d])[np.newaxis, np.newaxis, np.newaxis]) + scale0d = torch.from_numpy(np.array([self.scale0d])[np.newaxis, np.newaxis, np.newaxis]) + expected_transformed_image = (self.image + center0d) / self.scale0d + torch.testing.assert_close(transformed_image, expected_transformed_image) + received_image = affine_domain.receive(transformed_image) + expected_received_image = expected_transformed_image * scale0d - center0d + torch.testing.assert_close(received_image, expected_received_image) + + # when 1d + affine_domain = iamge_domain_module.AffineDomain(self.center1d, self.scale1d) + transformed_image = affine_domain.send(self.image) + center1d = self.center1d[np.newaxis, :, np.newaxis, np.newaxis] + scale1d = self.scale1d[np.newaxis, :, np.newaxis, np.newaxis] + expected_transformed_image = (self.image + center1d) / scale1d + torch.testing.assert_close(transformed_image, expected_transformed_image) + received_image = affine_domain.receive(transformed_image) + expected_received_image = expected_transformed_image * scale1d - center1d + torch.testing.assert_close(received_image, expected_received_image) + + # when 3d + affine_domain = iamge_domain_module.AffineDomain(self.center3d, self.scale3d) + transformed_image = affine_domain.send(self.image) + center3d = self.center3d[np.newaxis] + scale3d = self.scale3d[np.newaxis] + expected_transformed_image = (self.image + center3d) / scale3d + torch.testing.assert_close(transformed_image, expected_transformed_image) + received_image = affine_domain.receive(transformed_image) + expected_received_image = expected_transformed_image * scale3d - center3d + torch.testing.assert_close(received_image, expected_received_image) + + +class TestRGBDomain(unittest.TestCase): + """Tests fot bdpy.dl.torch.domain.image_domain.BGRDomain""" + + def setUp(self): + self.bgr_image = torch.rand((1, 3, 32, 32)) + self.rgb_image = self.bgr_image[:, [2, 1, 0], ...] + + def test_send(self): + """Test send""" + bgr_domain = iamge_domain_module.BGRDomain() + transformed_image = bgr_domain.send(self.bgr_image) + torch.testing.assert_close(transformed_image, self.rgb_image) + + def test_receive(self): + """Tests receive""" + bgr_domain = iamge_domain_module.BGRDomain() + received_image = bgr_domain.receive(self.rgb_image) + torch.testing.assert_close(received_image, self.bgr_image) + + +class TestPILDomainWithExplicitCrop(unittest.TestCase): + """Tests fot bdpy.dl.torch.domain.image_domain.PILDomainWithExplicitCrop""" + def setUp(self): + self.expected_transformed_image = torch.rand((1, 3, 32, 32)) + self.image = self.expected_transformed_image.permute(0, 2, 3, 1) * 255 + + def test_send(self): + """Test send""" + pdwe_domain = iamge_domain_module.PILDomainWithExplicitCrop() + transformed_image = pdwe_domain.send(self.image) + torch.testing.assert_close(transformed_image, self.expected_transformed_image) + + def test_receive(self): + """Tests receive""" + pdwe_domain = iamge_domain_module.PILDomainWithExplicitCrop() + with warnings.catch_warnings(record=True) as w: + received_image = pdwe_domain.receive(self.expected_transformed_image) + self.assertTrue(any(isinstance(warn.message, RuntimeWarning) for warn in w)) + torch.testing.assert_close(received_image, self.image) + + +class TestFixedResolutionDomain(unittest.TestCase): + """Tests fot bdpy.dl.torch.domain.image_domain.FixedResolutionDomain""" + def setUp(self): + self.expected_received_image_size = (1, 3, 16, 16) + self.image =torch.rand((1, 3, 32, 32)) + + def test_send(self): + """Test send""" + fr_domain = iamge_domain_module.FixedResolutionDomain((16, 16)) + with self.assertRaises(RuntimeError): + fr_domain.send(self.image) + + def test_receive(self): + """Tests receive""" + fr_domain = iamge_domain_module.FixedResolutionDomain((16, 16)) + + received_image = fr_domain.receive(self.image) + self.assertEqual(received_image.size(), self.expected_received_image_size) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/tests/recon/__init__.py b/tests/recon/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/recon/torch/__init__.py b/tests/recon/torch/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/recon/torch/modules/__init__.py b/tests/recon/torch/modules/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/recon/torch/modules/test_critic.py b/tests/recon/torch/modules/test_critic.py new file mode 100644 index 00000000..191b27f8 --- /dev/null +++ b/tests/recon/torch/modules/test_critic.py @@ -0,0 +1,173 @@ +"""Tests for bdpy.recon.torch.modules.critic.""" + +import unittest + +import torch + +from bdpy.recon.torch.modules import critic as critic_module + + +class TestBaseCritic(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.critic.BaseCritic.""" + def setUp(self): + self.features = { + "conv1": torch.tensor([1.0], requires_grad=True), + "conv2": torch.tensor([2.0], requires_grad=True), + "conv3": torch.tensor([3.0], requires_grad=True), + } + self.target_features = { + "conv1": torch.tensor([0.0]), + "conv2": torch.tensor([1.0]), + "conv3": torch.tensor([2.0]), + } + + def test_instantiation(self): + """Test instantiation.""" + self.assertRaises(TypeError, critic_module.BaseCritic) + + def test_call(self): + """Test __call__.""" + class ReturnZeroCritic(critic_module.BaseCritic): + def evaluate(self, features, target_features): + return 0.0 + + critic = ReturnZeroCritic() + self.assertEqual(critic(self.features, self.target_features), 0.0) + + def test_loss_computation(self): + """Test loss computation.""" + class SumCritic(critic_module.BaseCritic): + def evaluate(self, features, target_features): + loss = 0.0 + for feature, target_feature in zip(features.values(), target_features.values()): + loss += torch.sum(torch.abs(feature - target_feature)) + return loss + + critic = SumCritic() + self.assertEqual(critic(self.features, self.target_features), 3.0) + + for feature in self.features.values(): + feature.grad = None + loss = critic(self.features, self.target_features) + loss.backward() + for feature in self.features.values(): + self.assertIsNotNone(feature.grad) + self.assertEqual(feature.grad, torch.ones_like(feature)) + + +class TestNNModuleCritic(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.critic.NNModuleCritic.""" + def setUp(self): + self.features = { + "conv1": torch.tensor([1.0], requires_grad=True), + "conv2": torch.tensor([2.0], requires_grad=True), + "conv3": torch.tensor([3.0], requires_grad=True), + } + self.target_features = { + "conv1": torch.tensor([0.0]), + "conv2": torch.tensor([1.0]), + "conv3": torch.tensor([2.0]), + } + + def test_instantiation(self): + """Test instantiation.""" + self.assertRaises(TypeError, critic_module.NNModuleCritic) + + def test_call(self): + """Test __call__.""" + class ReturnZeroCritic(critic_module.NNModuleCritic): + def evaluate(self, features, target_features): + return 0.0 + + critic = ReturnZeroCritic() + self.assertEqual(critic(self.features, self.target_features), 0.0) + + def test_loss_computation(self): + """Test loss computation.""" + class SumCritic(critic_module.NNModuleCritic): + def evaluate(self, features, target_features): + loss = 0.0 + for feature, target_feature in zip(features.values(), target_features.values()): + loss += torch.sum(torch.abs(feature - target_feature)) + return loss + + critic = SumCritic() + self.assertEqual(critic(self.features, self.target_features), 3.0) + + for feature in self.features.values(): + feature.grad = None + loss = critic(self.features, self.target_features) + loss.backward() + for feature in self.features.values(): + self.assertIsNotNone(feature.grad) + self.assertEqual(feature.grad, torch.ones_like(feature)) + + +class TestLayerWiseAverageCritic(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.critic.LayerWiseAverageCritic.""" + def setUp(self): + self.features = { + "conv1": torch.tensor([1.0], requires_grad=True), + "conv2": torch.tensor([2.0], requires_grad=True), + "conv3": torch.tensor([3.0], requires_grad=True), + } + self.target_features = { + "conv1": torch.tensor([0.0]), + "conv2": torch.tensor([1.0]), + "conv3": torch.tensor([2.0]), + } + + def test_call(self): + """Test __call__.""" + class ReturnZeroCritic(critic_module.LayerWiseAverageCritic): + def compare_layer(self, feature, target_feature, layer_name): + return 0.0 + + critic = ReturnZeroCritic() + self.assertEqual(critic(self.features, self.target_features), 0.0) + + def test_loss_computation(self): + """Test loss computation.""" + class AbsCritic(critic_module.LayerWiseAverageCritic): + def compare_layer(self, feature, target_feature, layer_name): + return torch.abs(feature - target_feature) + + critic = AbsCritic() + self.assertEqual(critic(self.features, self.target_features), 1) + + for feature in self.features.values(): + feature.grad = None + loss = critic(self.features, self.target_features) + loss.backward() + for feature in self.features.values(): + self.assertIsNotNone(feature.grad) + self.assertEqual(feature.grad, torch.ones_like(feature)/3) + + +class TestMSE(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.critic.MSE.""" + def test_compare_layer(self): + """Test compare_layer.""" + critic = critic_module.MSE() + feature = torch.randn(13, 7) + target_feature = torch.randn_like(feature) + loss = critic.compare_layer(feature, target_feature, "conv1") + self.assertTrue(torch.allclose(loss, torch.sum((feature - target_feature)**2, dim=1))) + + +class TestTargetNormalizedMSE(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.critic.TargetNormalizedMSE.""" + def test_compare_layer(self): + """Test compare_layer.""" + critic = critic_module.TargetNormalizedMSE() + feature = torch.randn(13, 7) + target_feature = torch.randn_like(feature) + loss = critic.compare_layer(feature, target_feature, "conv1") + self.assertTrue(torch.allclose( + loss, + torch.sum((feature - target_feature)**2, dim=1)/torch.sum(target_feature**2, dim=1) + )) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/tests/recon/torch/modules/test_encoder.py b/tests/recon/torch/modules/test_encoder.py new file mode 100644 index 00000000..7d36ce8c --- /dev/null +++ b/tests/recon/torch/modules/test_encoder.py @@ -0,0 +1,120 @@ +"""Tests for bdpy.recon.torch.modules.encoder.""" + +import unittest + +import torch +import torch.nn as nn + +from bdpy.dl.torch.domain.image_domain import Zero2OneImageDomain +from bdpy.recon.torch.modules import encoder as encoder_module + + +class MLP(nn.Module): + """A simple MLP.""" + + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(64 * 64 * 3, 256) + self.fc2 = nn.Linear(256, 128) + + def forward(self, x): + x = x.view(x.size(0), -1) + x = self.fc1(x) + x = torch.relu(x) + x = self.fc2(x) + return x + + +class TestBaseEncoder(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.encoder.BaseEncoder.""" + + def test_instantiation(self): + """Test instantiation.""" + self.assertRaises(TypeError, encoder_module.BaseEncoder) + + def test_call(self): + """Test __call__.""" + + class ReturnAsIsEncoder(encoder_module.BaseEncoder): + def encode(self, images): + return {"image": images} + + encoder = ReturnAsIsEncoder() + images = torch.randn(1, 3, 64, 64) + features = encoder(images) + self.assertDictEqual(features, {"image": images}) + + +class TestNNModuleEncoder(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.encoder.NNModuleEncoder.""" + + def test_instantiation(self): + """Test instantiation.""" + self.assertRaises(TypeError, encoder_module.NNModuleEncoder) + + def test_call(self): + """Test __call__.""" + + class ReturnAsIsEncoder(encoder_module.NNModuleEncoder): + def __init__(self) -> None: + super().__init__() + def encode(self, images): + return {"image": images} + + encoder = ReturnAsIsEncoder() + + images = torch.randn(1, 3, 64, 64) + images.requires_grad = True + features = encoder(images) + self.assertIsInstance(features, dict) + self.assertEqual(len(features), 1) + self.assertEqual(features["image"].shape, (1, 3, 64, 64)) + features["image"].sum().backward() + self.assertIsNotNone(images.grad) + + +class TestSimpleEncoder(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.encoder.SimpleEncoder.""" + + def test_call(self): + """Test __call__.""" + encoder = encoder_module.SimpleEncoder( + MLP(), ["fc1", "fc2"], domain=Zero2OneImageDomain() + ) + images = torch.randn(1, 3, 64, 64).clamp(0, 1) + images.requires_grad = True + features = encoder(images) + self.assertIsInstance(features, dict) + self.assertEqual(len(features), 2) + self.assertEqual(features["fc1"].shape, (1, 256)) + self.assertEqual(features["fc2"].shape, (1, 128)) + features["fc2"].sum().backward() + self.assertIsNotNone(images.grad) + + +class TestBuildEncoder(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.encoder.build_encoder.""" + + def test_build_encoder(self): + """Test build_encoder.""" + mlp = MLP() + encoder_from_builder = encoder_module.build_encoder( + feature_network=mlp, + layer_names=["fc1", "fc2"], + domain=Zero2OneImageDomain(), + ) + encoder = encoder_module.SimpleEncoder( + mlp, ["fc1", "fc2"], domain=Zero2OneImageDomain() + ) + + images = torch.randn(1, 3, 64, 64).clamp(0, 1) + features_from_builder = encoder_from_builder(images) + features = encoder(images) + self.assertEqual(type(encoder_from_builder), type(encoder)) + self.assertEqual(features_from_builder.keys(), features.keys()) + for key in features_from_builder.keys(): + self.assertTrue(torch.allclose(features_from_builder[key], features[key])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/recon/torch/modules/test_generator.py b/tests/recon/torch/modules/test_generator.py new file mode 100644 index 00000000..965e391e --- /dev/null +++ b/tests/recon/torch/modules/test_generator.py @@ -0,0 +1,198 @@ +"""Tests for bdpy.recon.torch.modules.generator.""" + +import unittest + +import copy + +import torch +import torch.nn as nn +import torch.optim as optim +from torchvision.models import get_model + +from bdpy.recon.torch.modules import generator as generator_module + + +class LinearGenerator(generator_module.NNModuleGenerator): + def __init__(self): + super().__init__() + self.fc = nn.Linear(64, 10) + + def generate(self, latent): + return self.fc(latent) + + def reset_states(self) -> None: + self.fc.apply(generator_module.call_reset_parameters) + + +class TestCallResetParameters(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.generator.call_reset_parameters.""" + def setUp(self): + self.model_ids = [ + "alexnet", + "efficientnet_b0", + "fasterrcnn_resnet50_fpn", + "inception_v3", + "resnet18", + "vgg11", + "vit_b_16", + ] + # NOTE: The following modules are excluded from validation because they + # initialize their parameters as constants every time. + self.excluded_modules = [ + nn.modules.batchnorm._BatchNorm, + nn.LayerNorm, + ] + + def _validate_module(self, module: nn.Module, module_copy: nn.Module, parent_name: str = ""): + if isinstance(module, tuple(self.excluded_modules)): + return + for (name_p1, p1), (_, p2) in zip(module.named_parameters(recurse=False), module_copy.named_parameters(recurse=False)): + # NOTE: skip parameters that are prbably not randomly initialized + if "weight" not in name_p1: + continue + self.assertFalse( + torch.equal(p1, p2), + msg=f"Parameter {parent_name}.{name_p1} does not change after calling reset_parameters." + ) + for (name_m1, m1), (_, m2) in zip(module.named_children(), module_copy.named_children()): + self._validate_module(m1, m2, f"{parent_name}.{name_m1}") + + def test_call_reset_parameters(self): + """Test call_reset_parameters.""" + for model_id in self.model_ids: + model = get_model(model_id) + model_copy = copy.deepcopy(model) + for (name_p1, p1), (_, p2) in zip(model.named_parameters(), model_copy.named_parameters()): + self.assertTrue( + torch.equal(p1, p2), + msg=f"Parameter {name_p1} of {model_id} has been changed by deepcopy." + ) + model.apply(generator_module.call_reset_parameters) + self._validate_module(model, model_copy, model_id) + + +class TestBaseGenerator(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.generator.BaseGenerator.""" + + def test_instantiation(self): + """Test instantiation.""" + self.assertRaises(TypeError, generator_module.BaseGenerator) + + def test_call(self): + """Test __call__.""" + + class ReturnAsIsGenerator(generator_module.BaseGenerator): + def generate(self, latent): + return latent + + def reset_states(self) -> None: + pass + + def parameters(self, recurse=True): + return iter([]) + + generator = ReturnAsIsGenerator() + latent = torch.randn(1, 3, 64, 64) + generated_image = generator(latent) + self.assertEqual(generated_image.shape, (1, 3, 64, 64)) + + +class TestNNModuleGenerator(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.generator.NNModuleGenerator.""" + + def setUp(self): + """Set up.""" + self.generator = LinearGenerator() + + def test_instantiation(self): + """Test instantiation.""" + self.assertRaises(TypeError, generator_module.NNModuleGenerator) + + def test_call(self): + """Test __call__.""" + latent = torch.randn(1, 64) + generated_image = self.generator(latent) + self.assertEqual(generated_image.shape, (1, 10)) + generated_image.sum().backward() + self.assertIsNotNone(self.generator.fc.weight.grad) + + def test_reset_states(self): + """Test reset_states.""" + generator_copy = copy.deepcopy(self.generator) + for p1, p2 in zip(self.generator.parameters(), generator_copy.parameters()): + self.assertTrue(torch.equal(p1, p2)) + self.generator.reset_states() + for p1, p2 in zip(self.generator.parameters(), generator_copy.parameters()): + self.assertFalse(torch.equal(p1, p2)) + + +class TestBareGenerator(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.generator.BareGenerator.""" + + def test_call(self): + """Test __call__.""" + generator = generator_module.BareGenerator(activation=torch.sigmoid) + latent = torch.randn(1, 3, 64, 64) + generated_image = generator(latent) + self.assertEqual(generated_image.shape, (1, 3, 64, 64)) + torch.testing.assert_close(generated_image, torch.sigmoid(latent)) + + +class TestDNNGenerator(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.generator.DNNGenerator.""" + def test_call(self): + """Test __call__.""" + generator_network = LinearGenerator() + generator = generator_module.DNNGenerator(generator_network) + latent = torch.randn(1, 64) + generated_image = generator(latent) + self.assertEqual(generated_image.shape, (1, 10)) + generated_image.sum().backward() + self.assertIsNotNone(generator_network.fc.weight.grad) + + def test_reset_states(self): + """Test reset_states.""" + generator = generator_module.DNNGenerator(LinearGenerator()) + generator_copy = copy.deepcopy(generator) + for p1, p2 in zip(generator.parameters(), generator_copy.parameters()): + self.assertTrue(torch.equal(p1, p2)) + generator.reset_states() + for p1, p2 in zip(generator.parameters(), generator_copy.parameters()): + self.assertFalse(torch.equal(p1, p2)) + + +class TestFrozenGenerator(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.generator.FrozenGenerator.""" + def test_call(self): + """Test __call__.""" + generator_network = LinearGenerator() + generator = generator_module.FrozenGenerator(generator_network) + latent = torch.randn(1, 64) + generated_image = generator(latent) + self.assertEqual(generated_image.shape, (1, 10)) + self.assertRaises(ValueError, optim.SGD, generator.parameters()) + + def test_reset_states(self): + """Test reset_states.""" + generator = generator_module.FrozenGenerator(LinearGenerator()) + generator_copy = copy.deepcopy(generator) + for p1, p2 in zip(generator.parameters(), generator_copy.parameters()): + self.assertTrue(torch.equal(p1, p2)) + generator.reset_states() + for p1, p2 in zip(generator.parameters(), generator_copy.parameters()): + self.assertTrue(torch.equal(p1, p2)) + + +class TestBuildGenerator(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.generator.build_generator.""" + def test_build_generator(self): + """Test build_generator.""" + generator_network = LinearGenerator() + generator = generator_module.build_generator(generator_network) + self.assertIsInstance(generator, generator_module.DNNGenerator) + generator = generator_module.build_generator(generator_network, frozen=True) + self.assertIsInstance(generator, generator_module.FrozenGenerator) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/recon/torch/modules/test_latent.py b/tests/recon/torch/modules/test_latent.py new file mode 100644 index 00000000..6e24d50a --- /dev/null +++ b/tests/recon/torch/modules/test_latent.py @@ -0,0 +1,124 @@ +import torch +import unittest +from typing import Iterator +import torch.nn as nn +from functools import partial +from bdpy.recon.torch.modules import latent as latent_module + + +class DummyLatent(latent_module.BaseLatent): + def __init__(self): + self.latent = nn.Parameter(torch.tensor([0.0, 1.0, 2.0])) + + def reset_states(self): + with torch.no_grad(): + self.latent.fill_(0.0) + + def parameters(self, recurse): + return iter(self.latent) + + def generate(self): + return self.latent + + +class TestBaseLatent(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.latent.BaseLatent.""" + def setUp(self): + self.latent_value_expected = nn.Parameter(torch.tensor([0.0, 1.0, 2.0])) + self.latent_reset_value_expected = nn.Parameter(torch.tensor([0.0, 0.0, 0.0])) + + def test_instantiation(self): + """Test instantiation.""" + self.assertRaises(TypeError, latent_module.BaseLatent) + + def test_call(self): + """Test __call__.""" + latent = DummyLatent() + self.assertTrue(torch.equal(latent(), self.latent_value_expected)) + + def test_parameters(self): + """test parameters""" + latent = DummyLatent() + params = latent.parameters(recurse=True) + self.assertIsInstance(params, Iterator) + + def test_reset_states(self): + """test reset_states""" + latent = DummyLatent() + latent.reset_states() + self.assertTrue(torch.equal(latent(), self.latent_reset_value_expected)) + + +class DummyNNModuleLatent(latent_module.NNModuleLatent): + def __init__(self): + super().__init__() + self.latent = nn.Parameter(torch.tensor([0.0, 1.0, 2.0])) + + def reset_states(self): + with torch.no_grad(): + self.latent.fill_(0.0) + + def generate(self): + return self.latent + + +class TestNNModuleLatent(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.latent.NNModuleLatent.""" + def setUp(self): + self.latent_value_expected = nn.Parameter(torch.tensor([0.0, 1.0, 2.0])) + self.latent_reset_value_expected = nn.Parameter(torch.tensor([0.0, 0.0, 0.0])) + + def test_instantiation(self): + """Test instantiation.""" + self.assertRaises(TypeError, latent_module.NNModuleLatent) + + def test_call(self): + """Test __call__.""" + latent = DummyNNModuleLatent() + self.assertTrue(torch.equal(latent(), self.latent_value_expected)) + + def test_parameters(self): + """test parameters""" + latent = DummyNNModuleLatent() + params = latent.parameters(recurse=True) + self.assertIsInstance(params, Iterator) + + def test_reset_states(self): + """test reset_states""" + latent = DummyNNModuleLatent() + latent.reset_states() + self.assertTrue(torch.equal(latent(), self.latent_reset_value_expected)) + + +class TestArbitraryLatent(unittest.TestCase): + """Tests for bdpy.recon.torch.modules.latent.ArbitraryLatent.""" + def setUp(self): + self.latent = latent_module.ArbitraryLatent((1, 3, 64, 64), partial(nn.init.normal_, mean=0, std=1)) + self.latent_shape_expected = (1, 3, 64, 64) + self.latent_value_expected = nn.Parameter(torch.tensor([0.0, 1.0, 2.0])) + self.latent_reset_value_expected = nn.Parameter(torch.tensor([0.0, 0.0, 0.0])) + + def test_instantiation(self): + """Test instantiation.""" + self.assertRaises(TypeError, latent_module.ArbitraryLatent) + + def test_call(self): + """Test __call__.""" + self.assertEqual(self.latent().size(), self.latent_shape_expected) + + def test_parameters(self): + """test parameters""" + params = self.latent.parameters(recurse=True) + self.assertIsInstance(params, Iterator) + + def test_reset_states(self): + """test reset_states""" + self.latent.reset_states() + mean = self.latent().mean().item() + std = self.latent().std().item() + self.assertAlmostEqual(mean, 0, places=1) + self.assertAlmostEqual(std, 1, places=1) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/recon/torch/task/__init__.py b/tests/recon/torch/task/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/recon/torch/task/test_inversion.py b/tests/recon/torch/task/test_inversion.py new file mode 100644 index 00000000..080fe930 --- /dev/null +++ b/tests/recon/torch/task/test_inversion.py @@ -0,0 +1,223 @@ +"""Tests for bdpy.recon.torch.task.inversion""" + +from __future__ import annotations + +import unittest +from unittest.mock import patch, call +import copy +import torch +import torch.nn as nn +import torch.optim as optim + +from bdpy.recon.torch.task import inversion as inversion_module +from bdpy.task import callback as callback_module +from bdpy.dl.torch.domain.image_domain import Zero2OneImageDomain +from bdpy.recon.torch.modules import encoder as encoder_module +from bdpy.recon.torch.modules import generator as generator_module +from bdpy.recon.torch.modules import latent as latent_module +from bdpy.recon.torch.modules import critic as critic_module + + +class DummyFeatureInversionCallback(inversion_module.FeatureInversionCallback): + def __init__(self, total_steps = 1): + super().__init__() + self._total_steps = total_steps + self._loss = 0 + + def _step_str(self, step: int) -> str: + if self._total_steps > 0: + return f"{step+1}/{self._total_steps}" + else: + return f"{step+1}" + + def on_task_start(self): + print('task start') + + def on_iteration_start(self, step): + print(f"Step [{self._step_str(step)}] start") + + def on_image_generated(self, step, image): + print(f"Step [{self._step_str(step)}], {image.shape}") + + def on_loss_calculated(self, step, loss): + self._loss = loss.item() + + def on_iteration_end(self, step): + print(f"Step [{self._step_str(step)}] end") + + def on_task_end(self): + print('task end') + + + +class TestFeatureInversionCallback(unittest.TestCase): + """Tests for bdpy.recon.torch.task.inversion.FeatureInversionCallback""" + def setUp(self): + self.callback = inversion_module.FeatureInversionCallback() + self.expected_method_names = { + "on_task_start", + "on_iteration_start", + "on_image_generated", + "on_layerwise_loss_calculated", + "on_loss_calculated", + "on_iteration_end", + "on_task_end", + } + + def test_instance_methods(self): + method_names = { + event_type + for event_type in dir(self.callback) + if event_type.startswith("on_") + and callable(getattr(self.callback, event_type)) + } + self.assertEqual(method_names, self.expected_method_names) + for event_type in method_names: + fn = getattr(self.callback, event_type) + self.assertRaises(RuntimeError, fn) + + + def test_validate_callback(self): + + class Unrelated(callback_module.BaseCallback): + """Valid callback object but is not a subclass of TaskFeatureInversionCallback""" + + pass + + class HasUnknownEvent(DummyFeatureInversionCallback): + """Having invalid instance method `on_unknown_event` as a subclass of TaskFeatureInversionCallback""" + + def on_unknown_event(self): + pass + + self.assertIsNone( + callback_module._validate_callback(DummyFeatureInversionCallback(), inversion_module.FeatureInversionCallback) + ) + self.assertRaises( + TypeError, callback_module._validate_callback, Unrelated(), inversion_module.FeatureInversionCallback + ) + self.assertRaises(ValueError, HasUnknownEvent) + + +class TestCUILoggingCallback(unittest.TestCase): + """Tests for bdpy.recon.torch.task.inversion.CUILoggingCallback""" + def setUp(self): + self.callback = inversion_module.CUILoggingCallback() + self.expected_loss = torch.tensor([1.0]) + + def test_on_loss_culculated(self): + self.callback.on_loss_calculated(step=0, loss=self.expected_loss) + self.assertEqual(self.callback._loss, self.expected_loss.item()) + + @patch('builtins.print') + def test_on_iteration_end(self, mock_print): + self.callback.on_iteration_end(step=0) + mock_print.assert_called_once_with("Step: [1], Loss: -1.0000") + + +class MLP(nn.Module): + """A simple MLP.""" + + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(7 * 7 * 3, 32) + self.fc2 = nn.Linear(32, 10) + + def forward(self, x): + x = x.view(x.size(0), -1) + x = self.fc1(x) + x = torch.relu(x) + x = self.fc2(x) + return x + + +class LinearGenerator(generator_module.NNModuleGenerator): + def __init__(self): + super().__init__() + self.fc = nn.Linear(10, 7 * 7 * 3) + + def generate(self, latent): + return self.fc(latent) + + def reset_states(self) -> None: + self.fc.apply(generator_module.call_reset_parameters) + + +class DummyNNModuleLatent(latent_module.NNModuleLatent): + def __init__(self, base_latent): + super().__init__() + self.latent = nn.Parameter(base_latent) + + def reset_states(self): + with torch.no_grad(): + self.latent.fill_(0.0) + + def generate(self): + return self.latent + + +class TestFeatureInversionTask(unittest.TestCase): + """Tests for bdpy.recon.torch.task.inversion.FeatureInversionTask""" + def setUp(self): + self.init_latent = torch.randn(1, 10) + self.target_feature = { + 'fc1': torch.randn(1, 32), + 'fc2': torch.randn(1, 10) + } + self.encoder = encoder_module.SimpleEncoder( + MLP(), ["fc1", "fc2"], domain=Zero2OneImageDomain() + ) + self.generator = generator_module.DNNGenerator(LinearGenerator()) + self.latent = DummyNNModuleLatent(self.init_latent.clone()) + self.critic = critic_module.MSE() + self.optimizer = optim.SGD([self.latent.latent], lr=0.1) + self.callbacks = DummyFeatureInversionCallback() + + self.inversion_task = inversion_module.FeatureInversionTask( + encoder=self.encoder, + generator=self.generator, + latent=self.latent, + critic=self.critic, + optimizer=self.optimizer, + callbacks=self.callbacks + ) + + @patch('builtins.print') + def test_call(self, mock_print): + """Test __call__.""" + generated_image = self.inversion_task(self.target_feature) + self.assertTrue(len(self.inversion_task._callback_handler._callbacks) > 0) + + # test for process + assert isinstance(generated_image, torch.Tensor) + self.assertEqual(generated_image.shape, (1, 7 * 7 * 3)) + self.assertIsNotNone(self.inversion_task._generator._generator_network.fc.weight.grad) + self.assertFalse(torch.equal(self.inversion_task._latent.latent, self.init_latent)) + + + # test for callbacks + self.assertTrue(self.inversion_task._callback_handler._callbacks[0]._loss > 0 ) + mock_print.assert_has_calls([ + call('task start'), + call('Step [1/1] start'), + call('Step [1/1], torch.Size([1, 147])'), + call('Step [1/1] end'), + call('task end'), + ]) + + def test_reset_state(self): + """Test reset_states.""" + generator_copy = copy.deepcopy(self.inversion_task._generator) + latent_copy = copy.deepcopy(self.inversion_task._latent) + for p1, p2 in zip(self.inversion_task._generator.parameters(), generator_copy.parameters()): + self.assertTrue(torch.equal(p1, p2)) + torch.testing.assert_close(self.inversion_task._latent.latent, latent_copy.latent) + self.inversion_task.reset_states() + + for p1, p2 in zip(self.inversion_task._generator.parameters(), generator_copy.parameters()): + self.assertFalse(torch.equal(p1, p2)) + self.assertFalse(torch.equal(self.inversion_task._latent.latent, latent_copy.latent)) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/tests/task/__init__.py b/tests/task/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/task/test_callback.py b/tests/task/test_callback.py new file mode 100644 index 00000000..caae1cd2 --- /dev/null +++ b/tests/task/test_callback.py @@ -0,0 +1,218 @@ +"""Tests for bdpy.task.callback.""" + +from __future__ import annotations + +import unittest +from typing import Any, Callable + +from bdpy.task import callback as callback_module + + +# NOTE: setup functions +def setup_fns() -> list[tuple[Callable, tuple[Any], Any]]: + def f1(input_: Any) -> None: + pass + + def f2(input_): + pass + + def f3(a: int, b: int) -> int: + return a + b + + class F4: + def __call__(self, input_: Any) -> None: + pass + + return [ + (f1, (None,), None), + (f2, (None,), None), + (f3, (1, 2), 3), + (F4(), (None,), None), + ] + + +def setup_callback_classes(): + class TaskBaseCallback(callback_module.BaseCallback): + def __init__(self): + super().__init__(base_class=TaskBaseCallback) + + @callback_module.unused + def on_some_event(self, input_): + pass + + class AppendCallback(TaskBaseCallback): + def __init__(self): + self._storage = [] + + def on_some_event(self, input_): + self._storage.append(input_) + + return TaskBaseCallback, AppendCallback + + +class TestUnused(unittest.TestCase): + """Tests for unused decorator.""" + + def test_unused(self): + """Test unused decorator. + + Unused decorator should change the return type of the decorated function + to Annotated[None, "unused"]. The decorated function should raise a + RuntimeError when called. + + Examples + -------- + >>> @unused + ... def f(a: int, b: int, c: int = 0) -> int: + ... return a + b + c + ... + >>> f(1, 2, 3) + Traceback (most recent call last): + ... + RuntimeError: Function is decorated with @unused and must not be called. + """ + + params = setup_fns() + for fn, inputs_, output in params: + self.assertTrue( + not hasattr(fn, "__annotations__") + or fn.__annotations__.get("return", None) != callback_module._Unused + ) + self.assertEqual(fn(*inputs_), output) + unused_fn = callback_module.unused(fn) + self.assertTrue( + hasattr(unused_fn, "__annotations__") + and unused_fn.__annotations__.get("return", None) + == callback_module._Unused + ) + with self.assertRaises(RuntimeError): + unused_fn(*inputs_) + + def test_is_unused(self): + params = setup_fns() + for fn, _, _ in params: + self.assertFalse(callback_module._is_unused(fn)) + self.assertTrue(callback_module._is_unused(callback_module.unused(fn))) + + +class TestBaseCallback(unittest.TestCase): + def setUp(self): + self.callback = callback_module.BaseCallback() + self.expected_method_names = { + "on_task_start", + "on_task_end", + } + + def test_instance_methods(self): + method_names = { + event_type + for event_type in dir(self.callback) + if event_type.startswith("on_") + and callable(getattr(self.callback, event_type)) + } + self.assertEqual(method_names, self.expected_method_names) + for event_type in method_names: + fn = getattr(self.callback, event_type) + self.assertRaises(RuntimeError, fn) + + def test_subclass_definition(self): + TaskBaseCallback, _ = setup_callback_classes() + callback = TaskBaseCallback() + expected_method_names = {"on_task_start", "on_some_event", "on_task_end"} + method_names = { + event_type + for event_type in dir(callback) + if event_type.startswith("on_") and callable(getattr(callback, event_type)) + } + + self.assertEqual(method_names, expected_method_names) + for event_type in method_names: + fn = getattr(callback, event_type) + self.assertRaises(RuntimeError, fn) + + def test_validate_callback(self): + TaskBaseCallback, AppendCallback = setup_callback_classes() + + class Unrelated(callback_module.BaseCallback): + """Valid callback object but is not a subclass of TaskBaseCallback""" + + pass + + class HasUnknownEvent(TaskBaseCallback): + """Having invalid instance method `on_unknown_event` as a subclass of TaskBaseCallback""" + + def on_unknown_event(self): + pass + + self.assertIsNone( + callback_module._validate_callback(AppendCallback(), TaskBaseCallback) + ) + self.assertRaises( + TypeError, callback_module._validate_callback, Unrelated(), TaskBaseCallback + ) + self.assertRaises(ValueError, HasUnknownEvent) + + +class TestCallbackHandler(unittest.TestCase): + def test_initialization(self): + _, AppendCallback = setup_callback_classes() + c1, c2 = AppendCallback(), AppendCallback() + + handler = callback_module.CallbackHandler() + self.assertListEqual(handler._callbacks, []) + self.assertDictEqual(handler._registered_functions, {}) + + handler = callback_module.CallbackHandler(c1) + self.assertListEqual(handler._callbacks, [c1]) + self.assertDictEqual( + handler._registered_functions, + {"on_some_event": [c1.on_some_event]}, + ) + + handler = callback_module.CallbackHandler([c1, c2]) + self.assertListEqual(handler._callbacks, [c1, c2]) + self.assertDictEqual( + handler._registered_functions, + {"on_some_event": [c1.on_some_event, c2.on_some_event]}, + ) + + def test_register(self): + handler = callback_module.CallbackHandler() + _, AppendCallback = setup_callback_classes() + cb = AppendCallback() + + self.assertListEqual(handler._callbacks, []) + self.assertDictEqual(handler._registered_functions, {}) + handler.register(cb) + self.assertListEqual(handler._callbacks, [cb]) + self.assertDictEqual( + handler._registered_functions, + {"on_some_event": [cb.on_some_event]}, + ) + + def test_fire(self): + handler = callback_module.CallbackHandler() + _, AppendCallback = setup_callback_classes() + cb = AppendCallback() + handler.register(cb) + + self.assertListEqual(cb._storage, []) + + handler.fire("on_task_start") + self.assertListEqual(cb._storage, []) + + handler.fire("on_some_event", input_=1) + self.assertListEqual(cb._storage, [1]) + + handler.fire("on_some_event", input_=2) + self.assertListEqual(cb._storage, [1, 2]) + + # NOTE: fire() should only accept keyword arguments + self.assertRaises(TypeError, handler.fire, "on_some_event", 1, 2) + + handler.fire("on_task_end") + self.assertListEqual(cb._storage, [1, 2]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/task/test_core.py b/tests/task/test_core.py new file mode 100644 index 00000000..32fba222 --- /dev/null +++ b/tests/task/test_core.py @@ -0,0 +1,63 @@ +"""Tests for bdpy.task.core.""" + +from __future__ import annotations + +import unittest +from bdpy.task import core as core_module + + +class MockCallback(core_module.BaseCallback): + """Mock callback for testing.""" + def __init__(self): + self._storage = [] + + def on_some_event(self, input_): + self._storage.append(input_) + + +class MockTask(core_module.BaseTask[MockCallback]): + """Mock task for testing BaseTask.""" + def run(self, *inputs, **parameters): + self._callback_handler.fire("on_some_event", input_=1) + return inputs, parameters + + +class TestBaseTask(unittest.TestCase): + """Tests forbdpy.task.core.BaseTask """ + def setUp(self): + self.input1 = 1.0 + self.input2 = 2.0 + self.task_name = "reconstruction" + + def test_initialization_without_callbacks(self): + """Test initialization without callbacks.""" + task = MockTask() + self.assertIsInstance(task._callback_handler, core_module.CallbackHandler) + self.assertEqual(len(task._callback_handler._callbacks), 0) + + def test_initialization_with_callbacks(self): + """Test initialization with callbacks.""" + mock_callback = MockCallback() + task = MockTask(callbacks=mock_callback) + self.assertEqual(len(task._callback_handler._callbacks), 1) + self.assertIn(mock_callback, task._callback_handler._callbacks) + + def test_register_callback(self): + """Test register_callback method.""" + task = MockTask() + mock_callback = MockCallback() + task.register_callback(mock_callback) + self.assertIn(mock_callback, task._callback_handler._callbacks) + + def test_call(self): + """Test __call__""" + mock_callback = MockCallback() + task = MockTask(callbacks=mock_callback) + task_inputs, task_parameters = task(self.input1, self.input2, name=self.task_name) + self.assertEqual(task_inputs, (self.input1, self.input2)) + self.assertEqual(task_parameters["name"], self.task_name) + self.assertEqual(mock_callback._storage, [1]) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file