Сохранение / загрузка моделей PyTorch, завернутых в sklearn-совместимые оценщики

#python #scikit-learn #pickle #pytorch

#python #scikit-learn #рассол #pytorch

Вопрос:

Я написал автоэнкодер с использованием PyTorch и внедрил его в пользовательский базовый оценщик sklearn. Обычно я обучаю оценщик на машине с графическим процессором и сохраняю их для последующей оценки с помощью pickle. Если я пытаюсь загрузить оценщик на компьютер, где модель была сохранена на графическом процессоре, я получаю следующую ошибку:

Ошибка выполнения: попытка десериализации объекта на устройстве CUDA, но torch.cuda.is_available() имеет значение False. Если вы работаете на компьютере, использующем только процессор, пожалуйста, используйте torch.загрузите с помощью map_location=’cpu’, чтобы сопоставить ваши хранилища с процессором.

  • Есть ли способ принудительно переместить модель PyTorch в CPU перед обработкой оценщика без явного вызова?

  • Есть ли способ отменить выбор оценки, которая была сохранена, когда модель была на графическом процессоре?

Ниже приведен пример моей модели PyTorch и сопоставимого оценщика sklearn, а также пример того, как я пытаюсь сохранить и загрузить свои модели.

Модель PyTorch

 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 ) )
  

Оценщик, совместимый со sklearn

 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("self")
        for arg, val in values.items():
            setattr(self, arg, val)

        if use_cuda:
            if torch.cuda.is_available():
                self.device = torch.device("cuda")
            else:
                self.device = torch.device("cpu")
                warnings.warn("cuda not avaliable", UserWarning)
        else:
            self.device = torch.device("cpu")

    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 = "Train Epoch: {:5d},    Loss: {:15.6f}".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 = "Train Epoch: {:5d},    Loss: {:15.6f}".format(
                    epoch,
                    loss.item()
                )
                print message
        return model
  

Обучение

Обычно это делается на компьютере с графическим процессором.

 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)
  

Сохранение оценки

Я хотел бы найти способ модифицировать модель PyTorch в CPU до ее сохранения без явного вызова. Может быть, функция как часть класса AutoEncoder, которая вызывается по мере ее настройки?

 with open("autoencoder.pkl", "wb") 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)
  

Загрузка

Я не могу понять, как загрузить оценщик на машину без графического процессора, если он был сохранен, когда модель PyTorch все еще была на графическом процессоре.

 # This gives an error if the model was saved while on the GPU,
# and a GPU is not avaiable when loading.
with open("autoencoder.pkl", "rb") 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("autoencoder.pkl", "rb") as fp:
    touch.load(fp, map_location="cpu")
  

Комментарии:

1. Вам нужно удалить весь файл оценщика? Вы можете сохранить не только модель torch?

2. @ManojMohan Я бы хотел сохранить всю оценку, чтобы параметры обучения (эпохи, размер пакета и т.д.) были сохранены.).