Более быстрое вменение

Question

Я написал собственный объект Sci-kit Learn Transformer следующим образом:

from sklearn.base import BaseEstimator, TransformerMixin
from scipy.stats import norm, expon

#Custom transformer imputing 
class NumericalTransformer(BaseEstimator, TransformerMixin):

    y_labels = None
    x_labels = None
    X = None
    y = None
    df = None

    medians = dict()

    #Class Constructor
    def __init__(self, y_full):
        # y_full is a FULL LIST OF TARGET COLUMNS, BEFORE TRAIN/TEST/VALIDATION SPLIT
        # It is used only to merge targets to the given feature DataFrame in fit() and transform()
        #   methods, to carry out the median calculation/imputation logic. It it removed before
        #   the features are returned in transform() method, so no data leak occurs.
        self.y_full = y_full

    #helper function generating noise in imputed values
    def gen_noise(self, df, col):

        num_row = len(self.df)

        # generate noise of the same height as our dataset
        if ((df[col].dropna() >= 0) & (df[col].dropna() <= 1.0)).all():
            noise = 0
        elif (df[col].dropna() >= 0).all():
            noise = expon(loc = 0, scale = 0.5).rvs(num_row)
        else:
            noise = norm(loc = 0, scale = 1).rvs(num_row)

        # multiplication with isna() forces those at non-null values in df[col] to be 0
        return noise * self.df[col].isna()

    #Return self, nothing else to do here
    def fit(self, X, y):
        print('calling fit')

        NumericalTransformer.X = X.copy()
        NumericalTransformer.df = NumericalTransformer.X.merge(self.y_full, 
                                                           left_index = True, 
                                                           right_index = True, 
                                                           how = 'left')

        NumericalTransformer.y_labels = list(self.y_full.nunique().sort_values(ascending = False).index)
        NumericalTransformer.y_label = y.name
        NumericalTransformer.x_labels = [col for col in NumericalTransformer.X.select_dtypes(include = 'number').columns if col not in NumericalTransformer.y_labels]

        for tar in NumericalTransformer.y_labels:
            # medians collected from train indices only
            NumericalTransformer.medians[tar] = NumericalTransformer.df[NumericalTransformer.x_labels].groupby(by = NumericalTransformer.df[tar]).agg('median') 

        return self 

    #Custom transform method we wrote that creates aformentioned features and drops redundant ones 
    def transform(self, X, y = None):
        print('calling transform')

        NumericalTransformer.X = X.copy()
        NumericalTransformer.df = NumericalTransformer.X.merge(self.y_full, 
                                                           left_index = True, 
                                                           right_index = True, 
                                                           how = 'left')

        # Imputation of the missing values as per description above
        idx = NumericalTransformer.df.index

        for tar in NumericalTransformer.y_labels:
            NumericalTransformer.df.set_index(tar, inplace=True)

            for col in NumericalTransformer.x_labels:
                noise = gen_noise(NumericalTransformer.df, col)
                NumericalTransformer.df[col] = NumericalTransformer.df[col].fillna(NumericalTransformer.medians[tar][col]).add(noise.round(2) * pd.Series(data = NumericalTransformer.medians[tar][col].round(2), index = noise.index))

            NumericalTransformer.df.reset_index(inplace=True)

        NumericalTransformer.df.index = idx

        return NumericalTransformer.df[[col for col in df.columns if col not in NumericalTransformer.y_labels]]

Он предназначен для использования с такими данными:

df = pd.DataFrame(np.random.randint(0, 100, size=(vsize, 10)), 
  columns = ["col_{}".format(x) for x in range(10)], 
  index = range(0, vsize * 3, 3))

df_2 = pd.DataFrame(np.random.randint(0,100,size=(vsize, 10)), 
    columns = ["col_{}".format(x) for x in range(10, 20, 1)], 
    index = range(0, vsize * 2, 2))

df = df.merge(df_2, left_index = True, right_index = True, how = 'outer')

df_tar = pd.DataFrame({"tar_1": [np.random.randint(0, 2) for x in range(vsize * 3)], 
       "tar_2": [np.random.randint(0, 4) for x in range(vsize * 3)], 
       "tar_3": [np.random.randint(0, 8) for x in range(vsize * 3)], 
       "tar_4": [np.random.randint(0, 16) for x in range(vsize * 3)]})

df = df.merge(df_tar, left_index = True, right_index = True, how = 'inner')

Идея, лежащая в основе пользовательского объекта Transformer, заключается в том, чтобы инициировать полный список целей, а затем:

1) При установке в TRAIN set: запоминать значения MEDIAN на основе столбцов tar_1-4, в порядке убывания числа уникальных значений в них (например, tar_4: 16, tar_3: 8, tar_2: 4, tar_1: 2).

2) Когда преобразование используется в TRAIN set: заполните NaNs в каждом столбце значением MEDIAN, полученным из tar_4, или, если оно отсутствует: tar_3 / 2/1. Во всех случаях следует добавить некоторый шум, чтобы значения не были копиями других образцов подряд.

3) Когда преобразование используется в наборе ТЕСТ: то же, что в пункте 2

Объект создается с полным списком целей, так что не возникает проблем с методом transform (), когда .predict () вызывается в конвейере, частью которого он будет.

На данный момент метод fit () работает отлично, ему нужно всего ~ 120 мс, но transform () занял ~ 15 с с изящным ~ 1000 строками DataFrame.

Может кто-нибудь дать мне советы о том, как повысить его эффективность?