1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
|
# -*- coding: utf-8 -*-
"""
Utility functions for GPU
"""
# Author: Remi Flamary <remi.flamary@unice.fr>
# Nicolas Courty <ncourty@irisa.fr>
# Leo Gautheron <https://github.com/aje>
#
# License: MIT License
import cupy as np # np used for matrix computation
import cupy as cp # cp used for cupy specific operations
def euclidean_distances(a, b, squared=False, to_numpy=True):
"""
Compute the pairwise euclidean distance between matrices a and b.
Parameters
----------
a : np.ndarray (n, f)
first matrix
b : np.ndarray (m, f)
second matrix
gpu : boolean, optional (default False)
if True and the module cupy is available, the computation is done
on the GPU and the type of the matrix returned is cupy.ndarray.
Otherwise, compute on the CPU and returns np.ndarray.
squared : boolean, optional (default False)
if True, return squared euclidean distance matrix
Returns
-------
c : (n x m) np.ndarray or cupy.ndarray
pairwise euclidean distance distance matrix
"""
a, b = to_gpu(a, b)
a2=np.sum(np.square(a),1)
b2=np.sum(np.square(b),1)
c=-2*np.dot(a,b.T)
c+=a2[:,None]
c+=b2[None,:]
if not squared:
np.sqrt(c, out=c)
if to_numpy:
return to_np(c)
else:
return c
def dist(x1, x2=None, metric='sqeuclidean', to_numpy=True):
"""Compute distance between samples in x1 and x2 on gpu
Parameters
----------
x1 : np.array (n1,d)
matrix with n1 samples of size d
x2 : np.array (n2,d), optional
matrix with n2 samples of size d (if None then x2=x1)
metric : str
Metric from 'sqeuclidean', 'euclidean',
Returns
-------
M : np.array (n1,n2)
distance matrix computed with given metric
"""
if x2 is None:
x2 = x1
if metric == "sqeuclidean":
return euclidean_distances(x1, x2, squared=True, to_numpy=to_numpy)
elif metric == "euclidean":
return euclidean_distances(x1, x2, squared=False, to_numpy=to_numpy)
else:
raise NotImplementedError
def to_gpu(*args):
""" Upload numpy arrays to GPU and return them"""
if len(args) > 1:
return (cp.asarray(x) for x in args)
else:
return cp.asarray(args[0])
def to_np(*args):
""" convert GPU arras to numpy and return them"""
if len(args) > 1:
return (cp.asnumpy(x) for x in args)
else:
return cp.asnumpy(args[0])
|
