使用自定义距离指标保存 sklearn 模型

Question

我构建了一个带有自定义距离度量的 knn 模型，即余弦距离：

def cosine_distance(x,y):
    x_module = np.sqrt(np.sum(x**2))
    y_module = np.sqrt(np.sum(y**2))
    return 1-np.dot(x,y)/(x_module*y_module)

# load data
x_feature = load_npz('data/movie_features.npz').toarray()
movies = CSVHelper.read_movie('data/IMDB_Movies_Master_data.csv')

neigh = NearestNeighbors(n_neighbors=5, metric=cosine_distance)
neigh.fit(x_feature)

# save the k-means model
joblib.dump(neigh, 'knn.pkl')

现在在第二个脚本中，我使用以下命令加载模型joblib：

knn_classifier = joblib.load('knn.pkl')

但是，它会引发以下错误：

File "<stdin>", line 1, in <module>
  File "/home/A/deeplearning_env/lib/python3.5/site-packages/sklearn/externals/joblib/numpy_pickle.py", line 578, in load
    obj = _unpickle(fobj, filename, mmap_mode)
  File "/home/A/deeplearning_env/lib/python3.5/site-packages/sklearn/externals/joblib/numpy_pickle.py", line 508, in _unpickle
    obj = unpickler.load()
  File "/usr/lib/python3.5/pickle.py", line 1039, in load
    dispatch[key[0]](self)
  File "/usr/lib/python3.5/pickle.py", line 1334, in load_global
    klass = self.find_class(module, name)
  File "/usr/lib/python3.5/pickle.py", line 1388, in find_class
    return getattr(sys.modules[module], name)
AttributeError: module '__main__' has no attribute 'cosine_distance'

我如何知道joblib我正在使用自定义指标？我尝试cosine_distance在同一脚本中添加该函数，但它不起作用。