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
101
|
# -*- 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.
If the input matrix are in numpy format, they will be uploaded to the
GPU first which can incur significant time overhead.
Parameters
----------
a : np.ndarray (n, f)
first matrix
b : np.ndarray (m, f)
second matrix
to_numpy : boolean, optional (default True)
If true convert back the GPU array result to numpy format.
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])
|
