gpytorch models can be pickled by dill but not by cloudpickle

Open

#381 opened on Jun 16, 2020 2 comments 1 reaction 0 assignees View on GitHub

Dominant language: Python
Stars: 1.9k
Forks: 195
Avg merge: 1d 10h
Merged PRs (30d): 1

Description

Cloudpickle is failing at dumping [gpytorch](https://github.com/cornellius-gp/gpytorch) models. However they can be dumped with dill.

The following reproducible example
```python
import math
import torch
import gpytorch
import cloudpickle

train_x = torch.linspace(0, 1, 100)
train_y = torch.sin(train_x * (2 * math.pi)) + \
torch.randn(train_x.size()) * math.sqrt(0.04)

class ExactGPModel(gpytorch.models.ExactGP):
def __init__(self, train_x, train_y, likelihood):
super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
self.mean_module = gpytorch.means.ConstantMean()
self.covar_module = gpytorch.kernels.ScaleKernel(
gpytorch.kernels.RBFKernel())

def forward(self, x):
mean_x = self.mean_module(x)
covar_x = self.covar_module(x)
return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

# initialize likelihood and model
likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = ExactGPModel(train_x, train_y, likelihood)
cloudpickle.dumps(model)
```
returns the following error
```python-tb
~\AppData\Local\Continuum\anaconda3\envs\gp_env\lib\site-packages\torch\storage.py:34: FutureWarning: pickle support for Storage will be removed in 1.5. Use `torch.save` instead
warnings.warn("pickle support for Storage will be removed in 1.5. Use `torch.save` instead", FutureWarning)
---------------------------------------------------------------------------
PicklingError Traceback (most recent call last)
in
----> 1 pickled = cloudpickle.dumps(model)

~\AppData\Local\Continuum\anaconda3\envs\gp_env\lib\site-packages\cloudpickle\cloudpickle.py in dumps(obj, protocol)
1081 try:
1082 cp = CloudPickler(file, protocol=protocol)
-> 1083 cp.dump(obj)
1084 return file.getvalue()
1085 finally:

~\AppData\Local\Continuum\anaconda3\envs\gp_env\lib\site-packages\cloudpickle\cloudpickle.py in dump(self, obj)
475 self.inject_addons()
476 try:
--> 477 return Pickler.dump(self, obj)
478 except RuntimeError as e:
479 if 'recursion' in e.args[0]:

~\AppData\Local\Continuum\anaconda3\envs\gp_env\lib\pickle.py in dump(self, obj)
435 if self.proto >= 4:
436 self.framer.start_framing()
--> 437 self.save(obj)
438 self.write(STOP)
439 self.framer.end_framing()

~\AppData\Local\Continuum\anaconda3\envs\gp_env\lib\pickle.py in save(self, obj, save_persistent_id)
547
548 # Save the reduce() output and finally memoize the object
--> 549 self.save_reduce(obj=obj, *rv)
550
551 def persistent_id(self, obj):

~\AppData\Local\Continuum\anaconda3\envs\gp_env\lib\pickle.py in save_reduce(self, func, args, state, listitems, dictitems, obj)
660
661 if state is not None:
--> 662 save(state)
663 write(BUILD)
664

~\AppData\Local\Continuum\anaconda3\envs\gp_env\lib\pickle.py in save(self, obj, save_persistent_id)
502 f = self.dispatch.get(t)
503 if f is not None:
--> 504 f(self, obj) # Call unbound method with explicit self
505 return
506

~\AppData\Local\Continuum\anaconda3\envs\gp_env\lib\pickle.py in save_dict(self, obj)
854
855 self.memoize(obj)
--> 856 self._batch_setitems(obj.items())
857
858 dispatch[dict] = save_dict