title: Saving/Loading PyTorch Models Wrapped in sklearn-compatible Estimators
body: <p>I have written an Autoencoder using PyTorch and I have rolled it into a custom sklearn BaseEstimator. I normally train the estimator on a machine with a GPU and save them for later evaluation using pickle. If I try to load an estimator on a machine where the model was stored on the GPU, I get the following error:</p>
  <blockquote>
  <p>RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location='cpu' to map your storages to the CPU.</p>
  </blockquote>
  <ul>
  <li><p>Is there a way to force the PyTorch model to be moved to the CPU before pickling the estimator without an explicit call?</p>
  </li>
  <li><p>Is there a way to unpickle an estimator that was saved while the model was on the GPU?</p>
  </li>
  </ul>
  <p>The following is an example of my PyTorch model and sklearn comptatable estimator, along with an example of how I am trying to save and load my models.</p>
  <h1>PyTorch Model</h1>
  <pre><code>import torch.nn as nn
  class _AutoEncoder(nn.Module):
      def __init__(self, input_dim, output_dim, encoder_dim=4):
          super(_AutoEncoder, self).__init__()
          hidden_dim = int( (input_dim+encoder_dim)/2 )
          layers = []
          layers.append( nn.Linear(input_dim, hidden_dim) )
          layers.append( nn.Linear(hidden_dim, encoder_dim) )
          self.encoder = nn.Sequential(*layers)
          layers = []
          layers.append( nn.Linear(encoder_dim, hidden_dim) )
          layers.append( nn.Linear(hidden_dim, output_dim) )
          self.decoder = nn.Sequential(*layers)
  
      def forward(self, X):
          return self.decoder( self.encoder( X ) )
  </code></pre>
  <h1>sklearn Compatible Estimator</h1>
  <pre><code>import warnings
  import inspect
  import numpy as np
  from sklearn.base import BaseEstimator, TransformerMixin
  from sklearn import utils as sk_utils
  
  import torch
  import torch.nn as nn
  import torch.optim as optim
  import torch.utils.data
  
  class AutoEncoder(BaseEstimator, TransformerMixin):
      def __init__(
          self,
          encoder_dim=4,
          n_epochs=200,
          batch_size=None,
          shuffle=True,
          use_cuda=False
      ):
          super(AutoEncoder, self).__init__()
          args, _, _, values = inspect.getargvalues(inspect.currentframe())
          values.pop(&quot;self&quot;)
          for arg, val in values.items():
              setattr(self, arg, val)
  
          if use_cuda:
              if torch.cuda.is_available():
                  self.device = torch.device(&quot;cuda&quot;)
              else:
                  self.device = torch.device(&quot;cpu&quot;)
                  warnings.warn(&quot;cuda not avaliable&quot;, UserWarning)
          else:
              self.device = torch.device(&quot;cpu&quot;)
  
      def fit(self, X, y=None):
          # X, y = sk_utils.check_X_y(X, y, ensure_2d=False, allow_nd=True)
          self._model = self._train_classifier(X, y)
          return self
  
      def transform(self, X):
          sk_utils.validation.check_is_fitted(self, ['_model'])
          X = sk_utils.check_array(X)
          X = torch.from_numpy(X.astype(np.float32)).to(self.device)
          with torch.no_grad():
              self._model.eval()
              output = self._model.forward( X )
              return output.cpu().numpy()
  
      def encode(self, X):
          sk_utils.validation.check_is_fitted(self, ['_model'])
          X = sk_utils.check_array(X)
          X = torch.from_numpy(X.astype(np.float32)).to(self.device)
          with torch.no_grad():
              self._model.eval()
              output = self._model.encoder( X )
              return output.cpu().numpy()
  
      def decode(self, X):
          sk_utils.validation.check_is_fitted(self, ['_model'])
          X = sk_utils.check_array(X)
          X = torch.from_numpy(X.astype(np.float32)).to(self.device)
          with torch.no_grad():
              self._model.eval()
              output = self._model.decoder( X )
              return output.cpu().numpy()
  
      def _train_classifier(self, x_train, y_train):
          x_train = torch.from_numpy(x_train.astype(np.float32)).to(self.device)
          y_train = torch.from_numpy(y_train.astype(np.float32)).to(self.device)
          input_dim = x_train.shape[-1]
          output_dim = y_train.shape[-1]
          model = _AutoEncoder(input_dim, output_dim, encoder_dim=self.encoder_dim).to(self.device)
          loss_function = nn.MSELoss()
          optimizer = optim.Adam(model.parameters())
          print model
          if self.batch_size is None:
              return self._batch_train_simple_classifier(x_train, y_train, model, loss_function, optimizer)
          else:
              return self._minibatch_train_simple_classifier(x_train, y_train, model, loss_function, optimizer)
  
      def _batch_train_simple_classifier(self, x_train, y_train, model, loss_function, optimizer):
          for epoch in range(1, self.n_epochs+1):
              model.train()
              optimizer.zero_grad()
              outputs = model.forward(x_train)
              loss = loss_function(outputs, y_train)
              loss.backward()
              optimizer.step()
              if epoch % 10 == 0 or epoch == self.n_epochs:
                  message = &quot;Train Epoch: {:5d},    Loss: {:15.6f}&quot;.format(
                      epoch,
                      loss.item()
                  )
                  print message
          return model
  
      def _minibatch_train_simple_classifier(self, x_train, y_train, model, loss_function, optimizer):
          train_data = torch.utils.data.TensorDataset(x_train, y_train)
          train_loader = torch.utils.data.DataLoader(train_data, batch_size=self.batch_size, shuffle=self.shuffle)
          for epoch in range(1, self.n_epochs+1):
              for data, target in train_loader:
                  model.train()
                  optimizer.zero_grad()
                  outputs = model.forward(data)
                  loss = loss_function(outputs, target)
                  loss.backward()
                  optimizer.step()
              if epoch % 10 == 0 or epoch == self.n_epochs:
                  model.eval()
                  outputs = model.forward(x_train)
                  loss = loss_function(outputs, y_train)
                  message = &quot;Train Epoch: {:5d},    Loss: {:15.6f}&quot;.format(
                      epoch,
                      loss.item()
                  )
                  print message
          return model
  </code></pre>
  <h1>Training</h1>
  <p>This is normally done on a machine with a GPU.</p>
  <pre><code>from sklearn import datasets as sk_datasets
  digits = sk_datasets.load_digits(n_class=10, return_X_y=False)
  data = digits.data
  
  ae = AutoEncoder(
      encoder_dim=2,
      n_epochs=100,
      batch_size=128,
      shuffle=True,
      use_cuda=True
  )
  
  data_fitted = ae.fit_transform(data, data)
  </code></pre>
  <h1>Saving the Estimator</h1>
  <p>I'd like to find a way to have the PyTorch model moded to the CPU before it is saved without having an explicit call.  Maybe a function as part of the AutoEncoder class that gets called as it is pickled?</p>
  <pre><code>with open(&quot;autoencoder.pkl&quot;, &quot;wb&quot;) as fp:
      # ae._model needs to be moved to the CPU here.
      # I don't want to have to call ae._model.cpu() explicitly
      pickle.dump(ae, fp)
  </code></pre>
  <h1>Loading</h1>
  <p>I can not figure out how to load the estimator on a machine without a GPU if it was saved while the PyTorch model was still on the GPU.</p>
  <pre><code># This gives an error if the model was saved while on the GPU,
  # and a GPU is not avaiable when loading.
  with open(&quot;autoencoder.pkl&quot;, &quot;rb&quot;) as fp:
      model = pickle.load(fp)
  
  # This also a similar error. I also would not expect this to
  # work since the pickle file contains an sklearn estimator
  # wrapping a PyTorch model.
  with open(&quot;autoencoder.pkl&quot;, &quot;rb&quot;) as fp:
      touch.load(fp, map_location=&quot;cpu&quot;)
  </code></pre>