Source code for quaterion_models.model

from __future__ import annotations

import json
import os
from functools import partial
from typing import Any, Callable, Dict, List, Type, Union

import numpy as np
import torch
from torch import nn

from quaterion_models.encoders import Encoder
from quaterion_models.heads.encoder_head import EncoderHead
from quaterion_models.types import CollateFnType, MetaExtractorFnType, TensorInterchange
from quaterion_models.utils.classes import restore_class, save_class_import
from quaterion_models.utils.meta import merge_meta
from quaterion_models.utils.tensors import move_to_device

DEFAULT_ENCODER_KEY = "default"


[docs]class SimilarityModel(nn.Module):
    """Main class which contains encoder models with the head layer."""

    def __init__(self, encoders: Union[Encoder, Dict[str, Encoder]], head: EncoderHead):
        super().__init__()

        if not isinstance(encoders, dict):
            self.encoders: Dict[str, Encoder] = {DEFAULT_ENCODER_KEY: encoders}
        else:
            self.encoders: Dict[str, Encoder] = encoders

        for key, encoder in self.encoders.items():
            encoder.disable_gradients_if_required()
            self.add_module(key, encoder)

        self.head = head

[docs]    @classmethod
    def collate_fn(
        cls,
        batch: List[dict],
        encoders_collate_fns: Dict[str, CollateFnType],
        meta_extractors: Dict[str, MetaExtractorFnType],
    ) -> TensorInterchange:
        """Construct batches for all encoders

        Args:
            batch:
            encoders_collate_fns: Dict (or single) of collate functions associated with encoders
            meta_extractors: Dict (or single) of meta extractor functions associated with encoders

        """
        data = dict(
            (key, collate_fn(batch)) for key, collate_fn in encoders_collate_fns.items()
        )
        meta = dict(
            (key, meta_extractor_fn(batch))
            for key, meta_extractor_fn in meta_extractors.items()
        )
        return {
            "data": data,
            "meta": merge_meta(meta),
        }

[docs]    @classmethod
    def get_encoders_output_size(cls, encoders: Union[Encoder, Dict[str, Encoder]]):
        """Calculate total output size of given encoders

        Args:
            encoders:

        """
        encoders = encoders.values() if isinstance(encoders, dict) else [encoders]
        total_size = 0
        for encoder in encoders:
            total_size += encoder.embedding_size
        return total_size

[docs]    def train(self, mode: bool = True):
        super().train(mode)

[docs]    def get_collate_fn(self) -> Callable:
        """Construct a function to convert input data into neural network inputs

        Returns:
            neural network inputs
        """
        return partial(
            SimilarityModel.collate_fn,
            encoders_collate_fns=dict(
                (key, encoder.get_collate_fn())
                for key, encoder in self.encoders.items()
            ),
            meta_extractors=dict(
                (key, encoder.get_meta_extractor())
                for key, encoder in self.encoders.items()
            ),
        )

    # -------------------------------------------
    # ---------- Inference methods --------------
    # -------------------------------------------

[docs]    def encode(
        self, inputs: Union[List[Any], Any], batch_size=32, to_numpy=True
    ) -> Union[torch.Tensor, np.ndarray]:
        """Encode data in batches

        Args:
            inputs: list of input data to encode
            batch_size:
            to_numpy:

        Returns:
            Numpy array or torch.Tensor of shape (input_size, embedding_size)
        """
        self.eval()
        device = next(self.parameters(), torch.tensor(0)).device
        collate_fn = self.get_collate_fn()

        input_was_list = True
        if not isinstance(inputs, list):
            input_was_list = False
            inputs = [inputs]

        all_embeddings = []

        for start_index in range(0, len(inputs), batch_size):
            input_batch = [
                inputs[i]
                for i in range(start_index, min(len(inputs), start_index + batch_size))
            ]
            features = collate_fn(input_batch)
            features = move_to_device(features, device)

            with torch.no_grad():
                embeddings = self.forward(features)
                embeddings = embeddings.detach()
                if to_numpy:
                    embeddings = embeddings.cpu().numpy()
                all_embeddings.append(embeddings)

        if to_numpy:
            all_embeddings = np.concatenate(all_embeddings, axis=0)
        else:
            all_embeddings = torch.cat(all_embeddings, dim=0)

        if not input_was_list:
            all_embeddings = all_embeddings.squeeze()

        if to_numpy:
            all_embeddings = np.atleast_2d(all_embeddings)
        else:
            all_embeddings = torch.atleast_2d(all_embeddings)

        return all_embeddings

[docs]    def forward(self, batch):
        embeddings = [
            (key, encoder.forward(batch["data"][key]))
            for key, encoder in self.encoders.items()
        ]

        meta = batch["meta"]

        # Order embeddings by key name, to ensure reproduction
        embeddings = sorted(embeddings, key=lambda x: x[0])

        # Only embedding tensors of shape [batch_size x encoder_output_size]
        embedding_tensors = [embedding[1] for embedding in embeddings]

        # Shape: [batch_size x sum( encoders_emb_sizes )]
        joined_embeddings = torch.cat(embedding_tensors, dim=1)

        # Shape: [batch_size x output_emb_size]
        result_embedding = self.head(joined_embeddings, meta=meta)

        return result_embedding

    # -------------------------------------------
    # ---------- Persistence methods ------------
    # -------------------------------------------

    @classmethod
    def _get_head_path(cls, directory: str):
        return os.path.join(directory, "head")

    @classmethod
    def _get_encoders_path(cls, directory: str):
        return os.path.join(directory, "encoders")

[docs]    def save(self, output_path: str):
        head_path = self._get_head_path(output_path)
        os.makedirs(head_path, exist_ok=True)
        self.head.save(head_path)

        head_config = save_class_import(self.head)

        encoders_path = self._get_encoders_path(output_path)
        os.makedirs(encoders_path, exist_ok=True)

        encoders_config = []

        for encoder_key, encoder in self.encoders.items():
            encoder_path = os.path.join(encoders_path, encoder_key)
            os.mkdir(encoder_path)
            encoder.save(encoder_path)
            encoders_config.append({"key": encoder_key, **save_class_import(encoder)})

        with open(os.path.join(output_path, "config.json"), "w") as f_out:
            json.dump(
                {"encoders": encoders_config, "head": head_config}, f_out, indent=2
            )

[docs]    @classmethod
    def load(cls, input_path: str) -> SimilarityModel:
        with open(os.path.join(input_path, "config.json")) as f_in:
            config = json.load(f_in)

        head_config = config["head"]
        head_class: Type[EncoderHead] = restore_class(head_config)
        head_path = cls._get_head_path(input_path)
        head = head_class.load(head_path)

        encoders: Union[Encoder, Dict[str, Encoder]] = {}
        encoders_path = cls._get_encoders_path(input_path)
        encoders_config = config["encoders"]

        for encoder_params in encoders_config:
            encoder_key = encoder_params["key"]
            encoder_class = restore_class(encoder_params)
            encoders[encoder_key] = encoder_class.load(
                os.path.join(encoders_path, encoder_key)
            )

        return cls(head=head, encoders=encoders)


# In this framework, the terms Metric Learning and Similarity Learning are considered synonymous.
# However, the word "Metric" overlaps with other concepts in model training.
# In addition, the semantics of the word "Similarity" are simpler.
# It better reflects the basic idea of this training approach.
# That's why we prefer to use Similarity over Metric.
MetricModel = SimilarityModel