Source code for lumin.nn.models.blocks.endcap

import numpy as np
import pandas as pd
from typing import Union
from abc import abstractmethod

import torch.nn as nn
from torch import Tensor

from ....utils.misc import to_np

__all__ = ['AbsEndcap']

[docs]class AbsEndcap(nn.Module): r''' Abstract class for constructing post training layer which performs further calculation on NN outputs. Used when NN was trained to some proxy objective Arguments: model: trained :class:`~lumin.nn.models.model.Model` to wrap ''' def __init__(self, model:nn.Module): super().__init__() self.model = model
[docs] @abstractmethod def func(self, x:Tensor) -> Tensor: r''' Transformation functio to apply to model outputs Arguements: x: model output tensor Returns: Resulting tensor ''' pass
[docs] def forward(self, x:Tensor) -> Tensor: r''' Pass tensor through endcap and compute function Arguments: x: model output tensor Returns Resulting tensor ''' return self.func(x)
[docs] def predict(self, inputs:Union[np.ndarray, pd.DataFrame, Tensor], as_np:bool=True) -> Union[np.ndarray, Tensor]: r''' Evaluate model on input tensor, and comput function of model outputs Arguments: inputs: input data as Numpy array, Pandas DataFrame, or tensor on device as_np: whether to return predictions as Numpy array (otherwise tensor) Returns: model predictions pass through endcap function ''' # TODO add mask x = self.model.predict(inputs, as_np=False) x = self.func(x) return to_np(x) if as_np else x
Read the Docs v: v0.5.0
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.


Access comprehensive developer and user documentation for LUMIN

View Docs


Get tutorials for beginner and advanced researchers demonstrating many of the features of LUMIN

View Tutorials