Commit e70275cd authored by hazrmard's avatar hazrmard
Browse files

TorchEstimator defined, basic tests passing

parents
Pipeline #305 canceled with stages
*.ini
pytorchbridge.egg-info
# pyTorchBridge
`pytorchbridge` provides a `scikit-learn Estimator` API for pytorch modules:
```python
from pytorchbridge import TorchEstimator
import torch.nn as nn
import torch.optim as optim
# Define pyTorch module
mymodule = nn.Sequential(
nn.Conv2d(1,20,5),
nn.ReLU(),
nn.Conv2d(20,64,5),
nn.ReLU()
)
# Cast module to device (if not CPU)
mymodule.to(torch.device('cuda'))
# define optimizer, AFTER casting module to device
myoptim = optim.SGD(mymodule.parameters(), lr=0.1, momentum=0.9)
estimator = TorchEstimator(module=mymodule,
optimizer=myoptim,
loss=nn.MSELoss(),
epochs=5,
verbose=False,
batch_size=32,
cuda=True)
estimator.fit(X, y)
estimator.predict(X)
```
\ No newline at end of file
from .pytorchbridge import TorchEstimator
"""
Defines the `TorchEstimator` class which provides a Scikit-learn Estimator API
for pytorch modules.
"""
from typing import Iterable, Tuple, Iterator, Union
import numpy as np
from sklearn.base import BaseEstimator
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm.auto import trange, tqdm
class TorchEstimator(BaseEstimator):
"""
Wraps a `torch.nn.Module` instance with a scikit-learn `Estimator` API.
"""
def __init__(self, module: nn.Module=None,
optimizer: optim.Optimizer=None,
loss: nn.modules.loss._Loss=None, epochs: int=10, verbose=False,
batch_size: int=8, cuda=True):
"""
Keyword Arguments:
module {torch.nn.Module} -- A `nn.Module` describing the neural network,
optimizer {torch.optim.Optimizer} -- An `Optimizer` instance which
iteratively modifies weights,
loss {torch.nn._Loss} -- a `_Loss` instance which calculates the loss metric,
epochs {int} -- The number of times to iterate over the training data,
verbose {bool} -- Whether to log training progress or not,
batch_size {int} -- Chunk size of data for each training step,
cuda {bool} -- Whether to use GPU acceleration if available.
"""
self.module = module
self.optimizer = optimizer
self.loss = loss
self.epochs = epochs
self.verbose = verbose
self.batch_size = batch_size
self.cuda = cuda
# pylint: disable=no-member
self._device = torch.device('cpu')
self._batch_first = None
self._dtype = torch.float
def _init(self, X, y):
"""
Initializes internal parameters before fitting, including device, data
types for network parameters.
Arguments:
X {torch.Tensor} -- Features
y {torch.Tensor} -- Targets
"""
# pylint: disable=no-member
# Create a linear model if no module provided
if self.module is None:
self._device = torch.device('cuda') if \
torch.cuda.is_available() and self.cuda \
else torch.device('cpu')
_, _, infeatures = self._get_shape(X)
_, _, outfeatures = self._get_shape(y)
class MyModule(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(infeatures, outfeatures)
self.squeeze = len(torch.as_tensor(y).size()) == 1
def forward(self, x):
x = self.linear(x)
if self.squeeze:
return torch.squeeze(x)
return x
self.module = MyModule()
self.module.to(self._device)
else:
self._device = next(self.module.parameters()).device
self._dtype = next(self.module.parameters()).dtype
if self.optimizer is None:
self.optimizer = optim.SGD(self.module.parameters(), lr=0.1)
if self.loss is None:
self.loss = nn.MSELoss()
def fit(self, X: torch.Tensor, y: torch.Tensor) -> 'TorchEstimator':
"""
Fit target to features.
Arguments:
X {torch.Tensor} -- `Tensor` of shape (SeqLen, N, Features) or (N, SeqLen, Features)
for recurrent modules or (N, Features) for other modules.
y {torch.Tensor} -- `Tensor` of shape ([SeqLen,] N, OutputFeatures) for recurrent
modules of (N, OutputFeatures).
Returns:
self
"""
# pylint: disable=no-member
self._init(X, y)
if self.verbose:
print()
ranger = trange(self.epochs)
self._batch_first = self._is_batch_first()
for e in ranger:
total_loss = 0.
for instance, target in zip(self._to_batches(X), self._to_batches(y)):
instance, target = instance.to(self._device), target.to(self._device)
self.module.zero_grad()
output = self.module(instance)
loss = self.loss(output, target)
loss.backward()
self.optimizer.step()
total_loss += loss.item()
if self.verbose:
ranger.write(f'Epoch {e+1:3d}\tLoss: {total_loss:10.2f}')
return self
def predict(self, X: torch.Tensor) -> torch.Tensor:
"""
Predict output from inputs.
Arguments:
X {torch.Tensor} -- `Tensor` of shape (SeqLen, N, Features) or (N, SeqLen, Features)
for recurrent modules or (N, Features) for other modules.
Returns:
torch.Tensor -- of shape ([SeqLen,] N, OutputFeatures) for recurrent
modules of (N, OutputFeatures).
"""
# pylint: disable=no-member
is_numpy = isinstance(X, np.ndarray)
X = torch.as_tensor(X, dtype=self._dtype, device=self._device)
with torch.no_grad():
result = self.module(X)
if is_numpy:
return result.numpy()
return result
def score(self, X, y) -> float:
"""
Measure how well the estimator learned through the coefficient of
determination.
Arguments:
X {torch.Tensor} -- `Tensor` of shape (SeqLen, N, Features) or (N, SeqLen, Features)
for recurrent modules or (N, Features) for other modules.
y {torch.Tensor} -- `Tensor` of shape ([SeqLen,] N, OutputFeatures) for recurrent
modules of (N, OutputFeatures).
Returns:
float -- Coefficient of determination.
"""
y_pred = self.predict(X)
residual_squares_sum = ((y - y_pred) ** 2).sum()
total_squares_sum = ((y - y.mean()) ** 2).sum()
return (1 - residual_squares_sum / total_squares_sum).item()
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.
"""
# pylint: disable=no-member
if isinstance(X, np.ndarray):
X = X.astype(float)
X = torch.as_tensor(X, dtype=self._dtype)
if not self._batch_first:
# 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[0].size()[0]
N = len(X[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]
else:
# Fully connected layers take inputs of the shape (N, Features...)
N = len(X)
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:
"""
Checks whether the network has any recurrent units.
"""
return any(map(lambda x: isinstance(x, nn.RNNBase), self.module.modules()))
def _is_batch_first(self) -> bool:
"""
Checks whether the features arrays are in the shape (Batch, ..., Features) or
(..., Batch, Features).
"""
# Default setting is batch_first=False for RNNBase subclasses
return any(map(lambda x: getattr(x, 'batch_first', True), self.module.modules()))
def _get_shape(self, t: torch.Tensor) -> Tuple[int, int, int]:
"""
Get size of each dimension of tensor depending on `batch_first`. The
size is returned in order of time, batch, features.
Arguments:
t {torch.Tensor} -- A Tensor
Returns:
Tuple[int, int, int] -- A tuple of [time, batch, feature]
size of the tensor as interpreted by th estimator.
"""
# pylint: disable=no-member
if isinstance(t, np.ndarray):
sz = t.shape
else:
t = torch.as_tensor(t)
sz = t.size()
ndims = len(sz)
if ndims == 1:
return 0, sz[0], 1
if ndims == 2:
return 0, sz[0], sz[1]
elif ndims == 3:
if self._is_batch_first():
return sz[1], sz[0], sz[2]
else:
return sz[0], sz[1], sz[2]
from unittest import main, TestCase
import torch
import numpy as np
from sklearn.utils.estimator_checks import check_estimator
from sklearn.model_selection import GridSearchCV
from .pytorchbridge import TorchEstimator
class TestAPI(TestCase):
# def test_api(self):
# check_estimator(TorchEstimator())
def test_numpy_arrays(self):
X = np.random.rand(10, 2)
y = np.random.rand(10)
est = TorchEstimator()
est.fit(X, y).predict(X)
def test_torch_tensors(self):
# pylint: disable=no-member
X = torch.from_numpy(np.random.rand(10, 2))
y = torch.from_numpy(np.random.rand(10))
est = TorchEstimator()
est.fit(X, y).predict(X)
def grid_search(self):
pass
if __name__ == '__main__':
main()
from distutils.core import setup
setup(name='pytorchbridge',
version='0.1.0',
packages=['pytorchbridge'],
install_requires=['tqdm', 'scikit-learn>=0.20'],
author='Ibrahim Ahmed',
author_email='ibrahim.ahmed@vanderbilt.edu',
description='Scikit-learn Estimator API for pyTorch Modules',
url='https://git.isis.vanderbilt.edu/ahmedi/pyTorchBridge',
classifiers=[
'Development Status :: 3 - Alpha',
'Programming Language :: Python :: 3.6',
'Programming Language :: Python :: 3.7'
]
)
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment