finished TorchEstimator w/ fit, predict, score methods like scikit-learn Estimators

from multiprocessing import Pool

from os import cpu_count

from typing import Iterable, Tuple

from typing import Iterable, Tuple, Iterator

import numpy as np

import pandas as pd

class TorchEstimator:

    """

    Wraps a `torch.nn.Module` instance with a scikit-learn `Estimator` API.

    Args:

    * `module`: A `nn.Module` describing the neural network,

    * `optimizer`: An `Optimizer` instance which iteratively modifies weights,

    * `loss`: a `_Loss` instance which calculates the loss metric,

    * `epochs`: The number of times to iterate over the training data,

    * `verbose`: Whether to log training progress or not,

    * `batch_size`: Chunk size of data for each training step.

    """

    def __init__(self, module: nn.Module, optimizer: optim.Optimizer,

                 loss: nn.modules.loss._Loss, epochs: int=10, verbose=True):

                 loss: nn.modules.loss._Loss, epochs: int=10, verbose=True,

                 batch_size: int=8):

        self.module = module

        self.optimizer = optimizer

        self.loss = loss

        self.epochs = epochs

        self.verbose = verbose

        self.batch_size = batch_size

    def fit(self, X, y):

    def fit(self, X: torch.Tensor, y: torch.Tensor) -> 'TorchEstimator':

        """

        Fit target to features

        Fit target to features.

        Args:

        * `X`: `Tensor` of shape ([SeqLen,] N, Features) for recurrent modules or

        (N, Features). Mini-batches of shape (n, [SeqLen], Features) will be

        fed to the module at each iteration. So the module should re-view

        the tensors in the appropriate shape for the layers.

        * `y`: `Tensor` of shape ([SeqLen,] N, OutputFeatures) for recurrent

        modules of (N, OutputFeatures).

        Returns:

        * self

        """

        for _ in range(self.epochs):

            for instance, target in zip(self._to_batches(X), self._to_batches(y)):

        return self

    def predict(self, X):

        shape = X.size()

        results = []

    def predict(self, X: torch.Tensor) -> torch.Tensor:

        """

        Predict output from inputs.

        Args:

        * `X`: `Tensor` of shape ([SeqLen,] N, Features) for recurrent modules or

        (N, Features).

        Returns:

        * `Tensor` of shape ([SeqLen,] N, OutputFeatures) for recurrent

        modules of (N, OutputFeatures).

        """

        with torch.no_grad():

            for batch in self._to_batches(X):

                results.append(self.module(batch))

        res_tensor = torch.zeros(len(results), *results[-1].size(),

                                 dtype=results[-1].dtype)

        for i in range(len(res_tensor)):

            res_tensor[i] = results[i]

        return self._from_batches(res_tensor)

    def _to_batches(self, X):

        shape = X.size()

        ndims = len(shape)

        if ndims == 3:

            for i in range(shape[1]):

                yield X[:, i, :]

        else:

            for i in range(shape[0]):

                yield X[i]

            result = self.module(X)

        return result

    def score(self, X, y_true):

        """

        """

        y_pred = self.predict(X)

        residual_squares_sum = ((y_true - y_pred) ** 2).sum()

        total_squares_sum = ((y_true - y_true.mean()) ** 2).sum()

        return (1 - residual_squares_sum / total_squares_sum).item()

    def _from_batches(self, X):

        shape = X.size()

        ndims = len(shape)

        if ndims == 3:

            return X.transpose(0, 1)    # TODO: Transpose b/w seq len & batch not working

    def _to_batches(self, X: torch.Tensor) -> Iterator[torch.Tensor]:

        """

        Convert ([SeqLen,] N, Features) to a generator of ([SeqLen,] n, Features)

        mini-batches. So for recurrent layers, training can be done in batches.

        """

        if self._is_recurrent():

            # Recurrent layers take inputs of the shape (SeqLen, N, Features...)

            # So if there is any recurrent layer in the module, assume that this

            # is the expected input shape

            N = X.size()[1]

            nbatches = N // self.batch_size + (1 if N % self.batch_size else 0)

            for i in range(nbatches):

                yield X[:, i*self.batch_size:(i+1)*self.batch_size, :]

        else:

            return X

 No newline at end of file

            # Fully connected layers take inputsof the shape (N, Features...)

            N = X.size()[0]

            nbatches = N // self.batch_size + (1 if N % self.batch_size else 0)

            for i in range(nbatches):

                yield X[i*self.batch_size:(i+1)*self.batch_size]

    def _is_recurrent(self) -> bool:

        return any(map(lambda x: isinstance(x, nn.RNNBase), self.module.modules()))