Source code for nlpatl.models.embeddings.torchvision

from typing import List, Optional, Callable
import numpy as np

try:
    import torch
    import torchvision

    MODEL_FOR_TORCH_VISION_MAPPING_NAMES = {
        "resnet18": torchvision.models.resnet18,
        "alexnet": torchvision.models.alexnet,
        "vgg16": torchvision.models.vgg16,
    }
except ImportError:
    # No installation required if not using this function
    pass

from nlpatl.models.embeddings.embeddings import Embeddings


class TorchVision(Embeddings):
    """
    A wrapper of torch vision class.

    :param model_name_or_path: torch vision model name. Possible values
            are `resnet18`, `alexnet` and `vgg16`.
    :type model_name_or_path: str
    :param batch_size: Batch size of data processing. Default is 16
    :type batch_size: int
    :param model_config: Model paramateters. Refer to https://pytorch.org/vision/stable/models.html
    :type model_config: dict
    :param transform: Preprocessing function
    :type transform: :class: `torchvision.transforms`
    :param name: Name of this embeddings
    :type name: str

    >>> import nlpatl.models.embeddings as nme
    >>> model = nme.TrochVision()
    """

    def __init__(
        self,
        model_name_or_path: str,
        batch_size: int = 16,
        model_config: dict = {"pretrained": True},
        transform: Optional[Callable] = None,
        name: str = "torchvision",
    ):

        super().__init__(batch_size=batch_size, name=name)

        self.model_name_or_path = model_name_or_path
        self.model_config = model_config
        self.transform = transform

        if model_name_or_path in MODEL_FOR_TORCH_VISION_MAPPING_NAMES:
            self.model = MODEL_FOR_TORCH_VISION_MAPPING_NAMES[model_name_or_path](
                **model_config
            )
            self.model.eval()
        else:
            raise ValueError(
                "`{}` does not support. Supporting {} only".format(
                    model_name_or_path,
                    "`" + "`,`".join(MODEL_FOR_TORCH_VISION_MAPPING_NAMES.keys()) + "`",
                )
            )

    @staticmethod
    def get_mapping() -> dict:
        return MODEL_FOR_TORCH_VISION_MAPPING_NAMES

    def convert(self, x: List[np.ndarray]) -> np.ndarray:
        results = []
        for batch_inputs in self.batch(x, self.batch_size):
            with torch.no_grad():
                features = [
                    self.transform(img) if self.transform else img
                    for img in batch_inputs
                ]
            results.append(self.model(torch.stack(features)))

        return torch.cat(results).detach().numpy()