mars-project / mars

Mars is a tensor-based unified framework for large-scale data computation which scales numpy, pandas, scikit-learn and Python functions.
https://mars-project.readthedocs.io
Apache License 2.0
2.71k stars 327 forks source link

Does mars.learn support gpu -calculating ? #2019

Open LouisYZK opened 3 years ago

LouisYZK commented 3 years ago

I've tried transformed cuml's code to mars.lean in same api like bellow:

from mars.learn.neighbors import NearestNeighbors
def calculate_knn(emb: md.DataFrame, topk: int, device='cpu') -> mt.tensor:
    if device == 'gpu':
        emb = emb.to_gpu()
    nn = NearestNeighbors(n_neighbors=topk, metric='cosine')
    nn.fit(emb)
    d, i = nn.kneighbors(emb)
    return i

However, it seems it did not work on gpu for it's much slower than cuml in same scale data. Am I in a right way to lauch mars.learn's algorithm on gpu?

wjsi commented 3 years ago

Currently KNN does not have a GPU implementation.

LouisYZK commented 3 years ago

Currently KNN does not have a GPU implementation.

fine, thx.

qinxuye commented 3 years ago

This issue could keep open if this feature is required, we can keep on tracking via this issue.

LouisYZK commented 3 years ago

This issue could keep open if this feature is required, we can keep on tracking via this issue.

Thanks alot, we do need knn-calculating on GPU, for we do want mars to improve memory efficient instead of rapids/cuml.