mirror of
https://github.com/wassname/ray.git
synced 2026-07-19 11:27:32 +08:00
Switch build system to use CMake completely. (#200)
* switch to CMake completely ... * cleanup * Run C tests, update installation instructions.
This commit is contained in:
Philipp Moritz
committed by
Robert Nishihara
Robert Nishihara
parent
ba8933e10f
commit
a708e36225
@@ -0,0 +1,6 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
from .utils import copy_directory
|
||||
from .tfutils import TensorFlowVariables
|
||||
@@ -0,0 +1,7 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
from . import random
|
||||
from . import linalg
|
||||
from .core import *
|
||||
@@ -0,0 +1,227 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import ray.experimental.array.remote as ra
|
||||
import ray
|
||||
|
||||
__all__ = ["BLOCK_SIZE", "DistArray", "assemble", "zeros", "ones", "copy",
|
||||
"eye", "triu", "tril", "blockwise_dot", "dot", "transpose", "add", "subtract", "numpy_to_dist", "subblocks"]
|
||||
|
||||
BLOCK_SIZE = 10
|
||||
|
||||
class DistArray(object):
|
||||
def __init__(self, shape, objectids=None):
|
||||
self.shape = shape
|
||||
self.ndim = len(shape)
|
||||
self.num_blocks = [int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in self.shape]
|
||||
self.objectids = objectids if objectids is not None else np.empty(self.num_blocks, dtype=object)
|
||||
if self.num_blocks != list(self.objectids.shape):
|
||||
raise Exception("The fields `num_blocks` and `objectids` are inconsistent, `num_blocks` is {} and `objectids` has shape {}".format(self.num_blocks, list(self.objectids.shape)))
|
||||
|
||||
@staticmethod
|
||||
def compute_block_lower(index, shape):
|
||||
if len(index) != len(shape):
|
||||
raise Exception("The fields `index` and `shape` must have the same length, but `index` is {} and `shape` is {}.".format(index, shape))
|
||||
return [elem * BLOCK_SIZE for elem in index]
|
||||
|
||||
@staticmethod
|
||||
def compute_block_upper(index, shape):
|
||||
if len(index) != len(shape):
|
||||
raise Exception("The fields `index` and `shape` must have the same length, but `index` is {} and `shape` is {}.".format(index, shape))
|
||||
upper = []
|
||||
for i in range(len(shape)):
|
||||
upper.append(min((index[i] + 1) * BLOCK_SIZE, shape[i]))
|
||||
return upper
|
||||
|
||||
@staticmethod
|
||||
def compute_block_shape(index, shape):
|
||||
lower = DistArray.compute_block_lower(index, shape)
|
||||
upper = DistArray.compute_block_upper(index, shape)
|
||||
return [u - l for (l, u) in zip(lower, upper)]
|
||||
|
||||
@staticmethod
|
||||
def compute_num_blocks(shape):
|
||||
return [int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in shape]
|
||||
|
||||
def assemble(self):
|
||||
"""Assemble an array on this node from a distributed array of object IDs."""
|
||||
first_block = ray.get(self.objectids[(0,) * self.ndim])
|
||||
dtype = first_block.dtype
|
||||
result = np.zeros(self.shape, dtype=dtype)
|
||||
for index in np.ndindex(*self.num_blocks):
|
||||
lower = DistArray.compute_block_lower(index, self.shape)
|
||||
upper = DistArray.compute_block_upper(index, self.shape)
|
||||
result[[slice(l, u) for (l, u) in zip(lower, upper)]] = ray.get(self.objectids[index])
|
||||
return result
|
||||
|
||||
def __getitem__(self, sliced):
|
||||
# TODO(rkn): fix this, this is just a placeholder that should work but is inefficient
|
||||
a = self.assemble()
|
||||
return a[sliced]
|
||||
|
||||
# Register the DistArray class with Ray so that it knows how to serialize it.
|
||||
ray.register_class(DistArray)
|
||||
|
||||
@ray.remote
|
||||
def assemble(a):
|
||||
return a.assemble()
|
||||
|
||||
# TODO(rkn): what should we call this method
|
||||
@ray.remote
|
||||
def numpy_to_dist(a):
|
||||
result = DistArray(a.shape)
|
||||
for index in np.ndindex(*result.num_blocks):
|
||||
lower = DistArray.compute_block_lower(index, a.shape)
|
||||
upper = DistArray.compute_block_upper(index, a.shape)
|
||||
result.objectids[index] = ray.put(a[[slice(l, u) for (l, u) in zip(lower, upper)]])
|
||||
return result
|
||||
|
||||
@ray.remote
|
||||
def zeros(shape, dtype_name="float"):
|
||||
result = DistArray(shape)
|
||||
for index in np.ndindex(*result.num_blocks):
|
||||
result.objectids[index] = ra.zeros.remote(DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
|
||||
return result
|
||||
|
||||
@ray.remote
|
||||
def ones(shape, dtype_name="float"):
|
||||
result = DistArray(shape)
|
||||
for index in np.ndindex(*result.num_blocks):
|
||||
result.objectids[index] = ra.ones.remote(DistArray.compute_block_shape(index, shape), dtype_name=dtype_name)
|
||||
return result
|
||||
|
||||
@ray.remote
|
||||
def copy(a):
|
||||
result = DistArray(a.shape)
|
||||
for index in np.ndindex(*result.num_blocks):
|
||||
result.objectids[index] = a.objectids[index] # We don't need to actually copy the objects because cluster-level objects are assumed to be immutable.
|
||||
return result
|
||||
|
||||
@ray.remote
|
||||
def eye(dim1, dim2=-1, dtype_name="float"):
|
||||
dim2 = dim1 if dim2 == -1 else dim2
|
||||
shape = [dim1, dim2]
|
||||
result = DistArray(shape)
|
||||
for (i, j) in np.ndindex(*result.num_blocks):
|
||||
block_shape = DistArray.compute_block_shape([i, j], shape)
|
||||
if i == j:
|
||||
result.objectids[i, j] = ra.eye.remote(block_shape[0], block_shape[1], dtype_name=dtype_name)
|
||||
else:
|
||||
result.objectids[i, j] = ra.zeros.remote(block_shape, dtype_name=dtype_name)
|
||||
return result
|
||||
|
||||
@ray.remote
|
||||
def triu(a):
|
||||
if a.ndim != 2:
|
||||
raise Exception("Input must have 2 dimensions, but a.ndim is " + str(a.ndim))
|
||||
result = DistArray(a.shape)
|
||||
for (i, j) in np.ndindex(*result.num_blocks):
|
||||
if i < j:
|
||||
result.objectids[i, j] = ra.copy.remote(a.objectids[i, j])
|
||||
elif i == j:
|
||||
result.objectids[i, j] = ra.triu.remote(a.objectids[i, j])
|
||||
else:
|
||||
result.objectids[i, j] = ra.zeros_like.remote(a.objectids[i, j])
|
||||
return result
|
||||
|
||||
@ray.remote
|
||||
def tril(a):
|
||||
if a.ndim != 2:
|
||||
raise Exception("Input must have 2 dimensions, but a.ndim is " + str(a.ndim))
|
||||
result = DistArray(a.shape)
|
||||
for (i, j) in np.ndindex(*result.num_blocks):
|
||||
if i > j:
|
||||
result.objectids[i, j] = ra.copy.remote(a.objectids[i, j])
|
||||
elif i == j:
|
||||
result.objectids[i, j] = ra.tril.remote(a.objectids[i, j])
|
||||
else:
|
||||
result.objectids[i, j] = ra.zeros_like.remote(a.objectids[i, j])
|
||||
return result
|
||||
|
||||
@ray.remote
|
||||
def blockwise_dot(*matrices):
|
||||
n = len(matrices)
|
||||
if n % 2 != 0:
|
||||
raise Exception("blockwise_dot expects an even number of arguments, but len(matrices) is {}.".format(n))
|
||||
shape = (matrices[0].shape[0], matrices[n // 2].shape[1])
|
||||
result = np.zeros(shape)
|
||||
for i in range(n // 2):
|
||||
result += np.dot(matrices[i], matrices[n // 2 + i])
|
||||
return result
|
||||
|
||||
@ray.remote
|
||||
def dot(a, b):
|
||||
if a.ndim != 2:
|
||||
raise Exception("dot expects its arguments to be 2-dimensional, but a.ndim = {}.".format(a.ndim))
|
||||
if b.ndim != 2:
|
||||
raise Exception("dot expects its arguments to be 2-dimensional, but b.ndim = {}.".format(b.ndim))
|
||||
if a.shape[1] != b.shape[0]:
|
||||
raise Exception("dot expects a.shape[1] to equal b.shape[0], but a.shape = {} and b.shape = {}.".format(a.shape, b.shape))
|
||||
shape = [a.shape[0], b.shape[1]]
|
||||
result = DistArray(shape)
|
||||
for (i, j) in np.ndindex(*result.num_blocks):
|
||||
args = list(a.objectids[i, :]) + list(b.objectids[:, j])
|
||||
result.objectids[i, j] = blockwise_dot.remote(*args)
|
||||
return result
|
||||
|
||||
@ray.remote
|
||||
def subblocks(a, *ranges):
|
||||
"""
|
||||
This function produces a distributed array from a subset of the blocks in the `a`. The result and `a` will have the same number of dimensions.For example,
|
||||
subblocks(a, [0, 1], [2, 4])
|
||||
will produce a DistArray whose objectids are
|
||||
[[a.objectids[0, 2], a.objectids[0, 4]],
|
||||
[a.objectids[1, 2], a.objectids[1, 4]]]
|
||||
We allow the user to pass in an empty list [] to indicate the full range.
|
||||
"""
|
||||
ranges = list(ranges)
|
||||
if len(ranges) != a.ndim:
|
||||
raise Exception("sub_blocks expects to receive a number of ranges equal to a.ndim, but it received {} ranges and a.ndim = {}.".format(len(ranges), a.ndim))
|
||||
for i in range(len(ranges)):
|
||||
if ranges[i] == []: # We allow the user to pass in an empty list to indicate the full range
|
||||
ranges[i] = range(a.num_blocks[i])
|
||||
if not np.alltrue(ranges[i] == np.sort(ranges[i])):
|
||||
raise Exception("Ranges passed to sub_blocks must be sorted, but the {}th range is {}.".format(i, ranges[i]))
|
||||
if ranges[i][0] < 0:
|
||||
raise Exception("Values in the ranges passed to sub_blocks must be at least 0, but the {}th range is {}.".format(i, ranges[i]))
|
||||
if ranges[i][-1] >= a.num_blocks[i]:
|
||||
raise Exception("Values in the ranges passed to sub_blocks must be less than the relevant number of blocks, but the {}th range is {}, and a.num_blocks = {}.".format(i, ranges[i], a.num_blocks))
|
||||
last_index = [r[-1] for r in ranges]
|
||||
last_block_shape = DistArray.compute_block_shape(last_index, a.shape)
|
||||
shape = [(len(ranges[i]) - 1) * BLOCK_SIZE + last_block_shape[i] for i in range(a.ndim)]
|
||||
result = DistArray(shape)
|
||||
for index in np.ndindex(*result.num_blocks):
|
||||
result.objectids[index] = a.objectids[tuple([ranges[i][index[i]] for i in range(a.ndim)])]
|
||||
return result
|
||||
|
||||
@ray.remote
|
||||
def transpose(a):
|
||||
if a.ndim != 2:
|
||||
raise Exception("transpose expects its argument to be 2-dimensional, but a.ndim = {}, a.shape = {}.".format(a.ndim, a.shape))
|
||||
result = DistArray([a.shape[1], a.shape[0]])
|
||||
for i in range(result.num_blocks[0]):
|
||||
for j in range(result.num_blocks[1]):
|
||||
result.objectids[i, j] = ra.transpose.remote(a.objectids[j, i])
|
||||
return result
|
||||
|
||||
# TODO(rkn): support broadcasting?
|
||||
@ray.remote
|
||||
def add(x1, x2):
|
||||
if x1.shape != x2.shape:
|
||||
raise Exception("add expects arguments `x1` and `x2` to have the same shape, but x1.shape = {}, and x2.shape = {}.".format(x1.shape, x2.shape))
|
||||
result = DistArray(x1.shape)
|
||||
for index in np.ndindex(*result.num_blocks):
|
||||
result.objectids[index] = ra.add.remote(x1.objectids[index], x2.objectids[index])
|
||||
return result
|
||||
|
||||
# TODO(rkn): support broadcasting?
|
||||
@ray.remote
|
||||
def subtract(x1, x2):
|
||||
if x1.shape != x2.shape:
|
||||
raise Exception("subtract expects arguments `x1` and `x2` to have the same shape, but x1.shape = {}, and x2.shape = {}.".format(x1.shape, x2.shape))
|
||||
result = DistArray(x1.shape)
|
||||
for index in np.ndindex(*result.num_blocks):
|
||||
result.objectids[index] = ra.subtract.remote(x1.objectids[index], x2.objectids[index])
|
||||
return result
|
||||
@@ -0,0 +1,190 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import ray.experimental.array.remote as ra
|
||||
import ray
|
||||
|
||||
from .core import *
|
||||
|
||||
__all__ = ["tsqr", "modified_lu", "tsqr_hr", "qr"]
|
||||
|
||||
@ray.remote(num_return_vals=2)
|
||||
def tsqr(a):
|
||||
"""
|
||||
arguments:
|
||||
a: a distributed matrix
|
||||
Suppose that
|
||||
a.shape == (M, N)
|
||||
K == min(M, N)
|
||||
return values:
|
||||
q: DistArray, if q_full = ray.get(DistArray, q).assemble(), then
|
||||
q_full.shape == (M, K)
|
||||
np.allclose(np.dot(q_full.T, q_full), np.eye(K)) == True
|
||||
r: np.ndarray, if r_val = ray.get(np.ndarray, r), then
|
||||
r_val.shape == (K, N)
|
||||
np.allclose(r, np.triu(r)) == True
|
||||
"""
|
||||
if len(a.shape) != 2:
|
||||
raise Exception("tsqr requires len(a.shape) == 2, but a.shape is {}".format(a.shape))
|
||||
if a.num_blocks[1] != 1:
|
||||
raise Exception("tsqr requires a.num_blocks[1] == 1, but a.num_blocks is {}".format(a.num_blocks))
|
||||
|
||||
num_blocks = a.num_blocks[0]
|
||||
K = int(np.ceil(np.log2(num_blocks))) + 1
|
||||
q_tree = np.empty((num_blocks, K), dtype=object)
|
||||
current_rs = []
|
||||
for i in range(num_blocks):
|
||||
block = a.objectids[i, 0]
|
||||
q, r = ra.linalg.qr.remote(block)
|
||||
q_tree[i, 0] = q
|
||||
current_rs.append(r)
|
||||
for j in range(1, K):
|
||||
new_rs = []
|
||||
for i in range(int(np.ceil(1.0 * len(current_rs) / 2))):
|
||||
stacked_rs = ra.vstack.remote(*current_rs[(2 * i):(2 * i + 2)])
|
||||
q, r = ra.linalg.qr.remote(stacked_rs)
|
||||
q_tree[i, j] = q
|
||||
new_rs.append(r)
|
||||
current_rs = new_rs
|
||||
assert len(current_rs) == 1, "len(current_rs) = " + str(len(current_rs))
|
||||
|
||||
# handle the special case in which the whole DistArray "a" fits in one block
|
||||
# and has fewer rows than columns, this is a bit ugly so think about how to
|
||||
# remove it
|
||||
if a.shape[0] >= a.shape[1]:
|
||||
q_shape = a.shape
|
||||
else:
|
||||
q_shape = [a.shape[0], a.shape[0]]
|
||||
q_num_blocks = DistArray.compute_num_blocks(q_shape)
|
||||
q_objectids = np.empty(q_num_blocks, dtype=object)
|
||||
q_result = DistArray(q_shape, q_objectids)
|
||||
|
||||
# reconstruct output
|
||||
for i in range(num_blocks):
|
||||
q_block_current = q_tree[i, 0]
|
||||
ith_index = i
|
||||
for j in range(1, K):
|
||||
if np.mod(ith_index, 2) == 0:
|
||||
lower = [0, 0]
|
||||
upper = [a.shape[1], BLOCK_SIZE]
|
||||
else:
|
||||
lower = [a.shape[1], 0]
|
||||
upper = [2 * a.shape[1], BLOCK_SIZE]
|
||||
ith_index //= 2
|
||||
q_block_current = ra.dot.remote(q_block_current, ra.subarray.remote(q_tree[ith_index, j], lower, upper))
|
||||
q_result.objectids[i] = q_block_current
|
||||
r = current_rs[0]
|
||||
return q_result, ray.get(r)
|
||||
|
||||
# TODO(rkn): This is unoptimized, we really want a block version of this.
|
||||
@ray.remote(num_return_vals=3)
|
||||
def modified_lu(q):
|
||||
"""
|
||||
Algorithm 5 from http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf
|
||||
takes a matrix q with orthonormal columns, returns l, u, s such that q - s = l * u
|
||||
arguments:
|
||||
q: a two dimensional orthonormal q
|
||||
return values:
|
||||
l: lower triangular
|
||||
u: upper triangular
|
||||
s: a diagonal matrix represented by its diagonal
|
||||
"""
|
||||
q = q.assemble()
|
||||
m, b = q.shape[0], q.shape[1]
|
||||
S = np.zeros(b)
|
||||
|
||||
q_work = np.copy(q)
|
||||
|
||||
for i in range(b):
|
||||
S[i] = -1 * np.sign(q_work[i, i])
|
||||
q_work[i, i] -= S[i]
|
||||
q_work[(i + 1):m, i] /= q_work[i, i] # scale ith column of L by diagonal element
|
||||
q_work[(i + 1):m, (i + 1):b] -= np.outer(q_work[(i + 1):m, i], q_work[i, (i + 1):b]) # perform Schur complement update
|
||||
|
||||
L = np.tril(q_work)
|
||||
for i in range(b):
|
||||
L[i, i] = 1
|
||||
U = np.triu(q_work)[:b, :]
|
||||
return ray.get(numpy_to_dist.remote(ray.put(L))), U, S # TODO(rkn): get rid of put
|
||||
|
||||
@ray.remote(num_return_vals=2)
|
||||
def tsqr_hr_helper1(u, s, y_top_block, b):
|
||||
y_top = y_top_block[:b, :b]
|
||||
s_full = np.diag(s)
|
||||
t = -1 * np.dot(u, np.dot(s_full, np.linalg.inv(y_top).T))
|
||||
return t, y_top
|
||||
|
||||
@ray.remote
|
||||
def tsqr_hr_helper2(s, r_temp):
|
||||
s_full = np.diag(s)
|
||||
return np.dot(s_full, r_temp)
|
||||
|
||||
@ray.remote(num_return_vals=4)
|
||||
def tsqr_hr(a):
|
||||
"""Algorithm 6 from http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf"""
|
||||
q, r_temp = tsqr.remote(a)
|
||||
y, u, s = modified_lu.remote(q)
|
||||
y_blocked = ray.get(y)
|
||||
t, y_top = tsqr_hr_helper1.remote(u, s, y_blocked.objectids[0, 0], a.shape[1])
|
||||
r = tsqr_hr_helper2.remote(s, r_temp)
|
||||
return ray.get(y), ray.get(t), ray.get(y_top), ray.get(r)
|
||||
|
||||
@ray.remote
|
||||
def qr_helper1(a_rc, y_ri, t, W_c):
|
||||
return a_rc - np.dot(y_ri, np.dot(t.T, W_c))
|
||||
|
||||
@ray.remote
|
||||
def qr_helper2(y_ri, a_rc):
|
||||
return np.dot(y_ri.T, a_rc)
|
||||
|
||||
@ray.remote(num_return_vals=2)
|
||||
def qr(a):
|
||||
"""Algorithm 7 from http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf"""
|
||||
m, n = a.shape[0], a.shape[1]
|
||||
k = min(m, n)
|
||||
|
||||
# we will store our scratch work in a_work
|
||||
a_work = DistArray(a.shape, np.copy(a.objectids))
|
||||
|
||||
result_dtype = np.linalg.qr(ray.get(a.objectids[0, 0]))[0].dtype.name
|
||||
r_res = ray.get(zeros.remote([k, n], result_dtype)) # TODO(rkn): It would be preferable not to get this right after creating it.
|
||||
y_res = ray.get(zeros.remote([m, k], result_dtype)) # TODO(rkn): It would be preferable not to get this right after creating it.
|
||||
Ts = []
|
||||
|
||||
for i in range(min(a.num_blocks[0], a.num_blocks[1])): # this differs from the paper, which says "for i in range(a.num_blocks[1])", but that doesn't seem to make any sense when a.num_blocks[1] > a.num_blocks[0]
|
||||
sub_dist_array = subblocks.remote(a_work, list(range(i, a_work.num_blocks[0])), [i])
|
||||
y, t, _, R = tsqr_hr.remote(sub_dist_array)
|
||||
y_val = ray.get(y)
|
||||
|
||||
for j in range(i, a.num_blocks[0]):
|
||||
y_res.objectids[j, i] = y_val.objectids[j - i, 0]
|
||||
if a.shape[0] > a.shape[1]:
|
||||
# in this case, R needs to be square
|
||||
R_shape = ray.get(ra.shape.remote(R))
|
||||
eye_temp = ra.eye.remote(R_shape[1], R_shape[0], dtype_name=result_dtype)
|
||||
r_res.objectids[i, i] = ra.dot.remote(eye_temp, R)
|
||||
else:
|
||||
r_res.objectids[i, i] = R
|
||||
Ts.append(numpy_to_dist.remote(t))
|
||||
|
||||
for c in range(i + 1, a.num_blocks[1]):
|
||||
W_rcs = []
|
||||
for r in range(i, a.num_blocks[0]):
|
||||
y_ri = y_val.objectids[r - i, 0]
|
||||
W_rcs.append(qr_helper2.remote(y_ri, a_work.objectids[r, c]))
|
||||
W_c = ra.sum_list.remote(*W_rcs)
|
||||
for r in range(i, a.num_blocks[0]):
|
||||
y_ri = y_val.objectids[r - i, 0]
|
||||
A_rc = qr_helper1.remote(a_work.objectids[r, c], y_ri, t, W_c)
|
||||
a_work.objectids[r, c] = A_rc
|
||||
r_res.objectids[i, c] = a_work.objectids[i, c]
|
||||
|
||||
# construct q_res from Ys and Ts
|
||||
q = eye.remote(m, k, dtype_name=result_dtype)
|
||||
for i in range(len(Ts))[::-1]:
|
||||
y_col_block = subblocks.remote(y_res, [], [i])
|
||||
q = subtract.remote(q, dot.remote(y_col_block, dot.remote(Ts[i], dot.remote(transpose.remote(y_col_block), q))))
|
||||
|
||||
return ray.get(q), r_res
|
||||
@@ -0,0 +1,18 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import ray.experimental.array.remote as ra
|
||||
import ray
|
||||
|
||||
from .core import *
|
||||
|
||||
@ray.remote
|
||||
def normal(shape):
|
||||
num_blocks = DistArray.compute_num_blocks(shape)
|
||||
objectids = np.empty(num_blocks, dtype=object)
|
||||
for index in np.ndindex(*num_blocks):
|
||||
objectids[index] = ra.random.normal.remote(DistArray.compute_block_shape(index, shape))
|
||||
result = DistArray(shape, objectids)
|
||||
return result
|
||||
@@ -0,0 +1,7 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
from . import random
|
||||
from . import linalg
|
||||
from .core import *
|
||||
@@ -0,0 +1,86 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import ray
|
||||
|
||||
__all__ = ["zeros", "zeros_like", "ones", "eye", "dot", "vstack", "hstack", "subarray", "copy", "tril", "triu", "diag", "transpose", "add", "subtract", "sum", "shape", "sum_list"]
|
||||
|
||||
@ray.remote
|
||||
def zeros(shape, dtype_name="float", order="C"):
|
||||
return np.zeros(shape, dtype=np.dtype(dtype_name), order=order)
|
||||
|
||||
@ray.remote
|
||||
def zeros_like(a, dtype_name="None", order="K", subok=True):
|
||||
dtype_val = None if dtype_name == "None" else np.dtype(dtype_name)
|
||||
return np.zeros_like(a, dtype=dtype_val, order=order, subok=subok)
|
||||
|
||||
@ray.remote
|
||||
def ones(shape, dtype_name="float", order="C"):
|
||||
return np.ones(shape, dtype=np.dtype(dtype_name), order=order)
|
||||
|
||||
@ray.remote
|
||||
def eye(N, M=-1, k=0, dtype_name="float"):
|
||||
M = N if M == -1 else M
|
||||
return np.eye(N, M=M, k=k, dtype=np.dtype(dtype_name))
|
||||
|
||||
@ray.remote
|
||||
def dot(a, b):
|
||||
return np.dot(a, b)
|
||||
|
||||
@ray.remote
|
||||
def vstack(*xs):
|
||||
return np.vstack(xs)
|
||||
|
||||
@ray.remote
|
||||
def hstack(*xs):
|
||||
return np.hstack(xs)
|
||||
|
||||
# TODO(rkn): instead of this, consider implementing slicing
|
||||
@ray.remote
|
||||
def subarray(a, lower_indices, upper_indices): # TODO(rkn): be consistent about using "index" versus "indices"
|
||||
return a[[slice(l, u) for (l, u) in zip(lower_indices, upper_indices)]]
|
||||
|
||||
@ray.remote
|
||||
def copy(a, order="K"):
|
||||
return np.copy(a, order=order)
|
||||
|
||||
@ray.remote
|
||||
def tril(m, k=0):
|
||||
return np.tril(m, k=k)
|
||||
|
||||
@ray.remote
|
||||
def triu(m, k=0):
|
||||
return np.triu(m, k=k)
|
||||
|
||||
@ray.remote
|
||||
def diag(v, k=0):
|
||||
return np.diag(v, k=k)
|
||||
|
||||
@ray.remote
|
||||
def transpose(a, axes=[]):
|
||||
axes = None if axes == [] else axes
|
||||
return np.transpose(a, axes=axes)
|
||||
|
||||
@ray.remote
|
||||
def add(x1, x2):
|
||||
return np.add(x1, x2)
|
||||
|
||||
@ray.remote
|
||||
def subtract(x1, x2):
|
||||
return np.subtract(x1, x2)
|
||||
|
||||
@ray.remote
|
||||
def sum(x, axis=-1):
|
||||
return np.sum(x, axis=axis if axis != -1 else None)
|
||||
|
||||
@ray.remote
|
||||
def shape(a):
|
||||
return np.shape(a)
|
||||
|
||||
# We use Any to allow different numerical types as well as numpy arrays.
|
||||
# TODO(rkn):this isn't in the numpy API, so be careful about exposing this.
|
||||
@ray.remote
|
||||
def sum_list(*xs):
|
||||
return np.sum(xs, axis=0)
|
||||
@@ -0,0 +1,91 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import ray
|
||||
|
||||
__all__ = ["matrix_power", "solve", "tensorsolve", "tensorinv", "inv",
|
||||
"cholesky", "eigvals", "eigvalsh", "pinv", "slogdet", "det",
|
||||
"svd", "eig", "eigh", "lstsq", "norm", "qr", "cond", "matrix_rank",
|
||||
"LinAlgError", "multi_dot"]
|
||||
|
||||
@ray.remote
|
||||
def matrix_power(M, n):
|
||||
return np.linalg.matrix_power(M, n)
|
||||
|
||||
@ray.remote
|
||||
def solve(a, b):
|
||||
return np.linalg.solve(a, b)
|
||||
|
||||
@ray.remote(num_return_vals=2)
|
||||
def tensorsolve(a):
|
||||
raise NotImplementedError
|
||||
|
||||
@ray.remote(num_return_vals=2)
|
||||
def tensorinv(a):
|
||||
raise NotImplementedError
|
||||
|
||||
@ray.remote
|
||||
def inv(a):
|
||||
return np.linalg.inv(a)
|
||||
|
||||
@ray.remote
|
||||
def cholesky(a):
|
||||
return np.linalg.cholesky(a)
|
||||
|
||||
@ray.remote
|
||||
def eigvals(a):
|
||||
return np.linalg.eigvals(a)
|
||||
|
||||
@ray.remote
|
||||
def eigvalsh(a):
|
||||
raise NotImplementedError
|
||||
|
||||
@ray.remote
|
||||
def pinv(a):
|
||||
return np.linalg.pinv(a)
|
||||
|
||||
@ray.remote
|
||||
def slogdet(a):
|
||||
raise NotImplementedError
|
||||
|
||||
@ray.remote
|
||||
def det(a):
|
||||
return np.linalg.det(a)
|
||||
|
||||
@ray.remote(num_return_vals=3)
|
||||
def svd(a):
|
||||
return np.linalg.svd(a)
|
||||
|
||||
@ray.remote(num_return_vals=2)
|
||||
def eig(a):
|
||||
return np.linalg.eig(a)
|
||||
|
||||
@ray.remote(num_return_vals=2)
|
||||
def eigh(a):
|
||||
return np.linalg.eigh(a)
|
||||
|
||||
@ray.remote(num_return_vals=4)
|
||||
def lstsq(a, b):
|
||||
return np.linalg.lstsq(a)
|
||||
|
||||
@ray.remote
|
||||
def norm(x):
|
||||
return np.linalg.norm(x)
|
||||
|
||||
@ray.remote(num_return_vals=2)
|
||||
def qr(a):
|
||||
return np.linalg.qr(a)
|
||||
|
||||
@ray.remote
|
||||
def cond(x):
|
||||
return np.linalg.cond(x)
|
||||
|
||||
@ray.remote
|
||||
def matrix_rank(M):
|
||||
return np.linalg.matrix_rank(M)
|
||||
|
||||
@ray.remote
|
||||
def multi_dot(*a):
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,10 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
import ray
|
||||
|
||||
@ray.remote
|
||||
def normal(shape):
|
||||
return np.random.normal(size=shape)
|
||||
@@ -0,0 +1,78 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
import numpy as np
|
||||
|
||||
def unflatten(vector, shapes):
|
||||
i = 0
|
||||
arrays = []
|
||||
for shape in shapes:
|
||||
size = np.prod(shape)
|
||||
array = vector[i:(i + size)].reshape(shape)
|
||||
arrays.append(array)
|
||||
i += size
|
||||
assert len(vector) == i, "Passed weight does not have the correct shape."
|
||||
return arrays
|
||||
|
||||
class TensorFlowVariables(object):
|
||||
"""An object used to extract variables from a loss function.
|
||||
|
||||
This object also provides methods for getting and setting the weights of the
|
||||
relevant variables.
|
||||
|
||||
Attributes:
|
||||
sess (tf.Session): The tensorflow session used to run assignment.
|
||||
loss: The loss function passed in by the user.
|
||||
variables (List[tf.Variable]): Extracted variables from the loss.
|
||||
assignment_placeholders (List[tf.placeholders]): The nodes that weights get
|
||||
passed to.
|
||||
assignment_nodes (List[tf.Tensor]): The nodes that assign the weights.
|
||||
"""
|
||||
def __init__(self, loss, sess=None):
|
||||
"""Creates a TensorFlowVariables instance."""
|
||||
import tensorflow as tf
|
||||
self.sess = sess
|
||||
self.loss = loss
|
||||
variable_names = [op.node_def.name for op in loss.graph.get_operations() if op.node_def.op == "Variable"]
|
||||
self.variables = [v for v in tf.trainable_variables() if v.op.node_def.name in variable_names]
|
||||
self.assignment_placeholders = dict()
|
||||
self.assignment_nodes = []
|
||||
|
||||
# Create new placeholders to put in custom weights.
|
||||
for var in self.variables:
|
||||
self.assignment_placeholders[var.op.node_def.name] = tf.placeholder(var.value().dtype, var.get_shape().as_list())
|
||||
self.assignment_nodes.append(var.assign(self.assignment_placeholders[var.op.node_def.name]))
|
||||
|
||||
def set_session(self, sess):
|
||||
"""Modifies the current session used by the class."""
|
||||
self.sess = sess
|
||||
|
||||
def get_flat_size(self):
|
||||
return sum([np.prod(v.get_shape().as_list()) for v in self.variables])
|
||||
|
||||
def _check_sess(self):
|
||||
"""Checks if the session is set, and if not throw an error message."""
|
||||
assert self.sess is not None, "The session is not set. Set the session either by passing it into the TensorFlowVariables constructor or by calling set_session(sess)."
|
||||
|
||||
def get_flat(self):
|
||||
"""Gets the weights and returns them as a flat array."""
|
||||
self._check_sess()
|
||||
return np.concatenate([v.eval(session=self.sess).flatten() for v in self.variables])
|
||||
|
||||
def set_flat(self, new_weights):
|
||||
"""Sets the weights to new_weights, converting from a flat array."""
|
||||
self._check_sess()
|
||||
shapes = [v.get_shape().as_list() for v in self.variables]
|
||||
arrays = unflatten(new_weights, shapes)
|
||||
placeholders = [self.assignment_placeholders[v.op.node_def.name] for v in self.variables]
|
||||
self.sess.run(self.assignment_nodes, feed_dict=dict(zip(placeholders,arrays)))
|
||||
|
||||
def get_weights(self):
|
||||
"""Returns the weights of the variables of the loss function in a list."""
|
||||
self._check_sess()
|
||||
return {v.op.node_def.name: v.eval(session=self.sess) for v in self.variables}
|
||||
|
||||
def set_weights(self, new_weights):
|
||||
"""Sets the weights to new_weights."""
|
||||
self._check_sess()
|
||||
self.sess.run(self.assignment_nodes, feed_dict={self.assignment_placeholders[name]: value for (name, value) in new_weights.items()})
|
||||
@@ -0,0 +1,74 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import io
|
||||
import os
|
||||
import tarfile
|
||||
import sys
|
||||
|
||||
import ray
|
||||
|
||||
def tarred_directory_as_bytes(source_dir):
|
||||
"""Tar a directory and return it as a byte string.
|
||||
|
||||
Args:
|
||||
source_dir (str): The name of the directory to tar.
|
||||
|
||||
Returns:
|
||||
A byte string representing the tarred file.
|
||||
"""
|
||||
# Get a BytesIO object.
|
||||
string_file = io.BytesIO()
|
||||
# Create an in-memory tarfile of the source directory.
|
||||
with tarfile.open(mode="w:gz", fileobj=string_file) as tar:
|
||||
tar.add(source_dir, arcname=os.path.basename(source_dir))
|
||||
string_file.seek(0)
|
||||
return string_file.read()
|
||||
|
||||
def tarred_bytes_to_directory(tarred_bytes, target_dir):
|
||||
"""Take a byte string and untar it.
|
||||
|
||||
Args:
|
||||
tarred_bytes (str): A byte string representing the tarred file. This should
|
||||
be the output of tarred_directory_as_bytes.
|
||||
target_dir (str): The directory to create the untarred files in.
|
||||
"""
|
||||
string_file = io.BytesIO(tarred_bytes)
|
||||
with tarfile.open(fileobj=string_file) as tar:
|
||||
tar.extractall(path=target_dir)
|
||||
|
||||
def copy_directory(source_dir, target_dir=None):
|
||||
"""Copy a local directory to each machine in the Ray cluster.
|
||||
|
||||
Note that both source_dir and target_dir must have the same basename). For
|
||||
example, source_dir can be /a/b/c and target_dir can be /d/e/c. In this case,
|
||||
the directory /d/e will be added to the Python path of each worker.
|
||||
|
||||
Note that this method is not completely safe to use. For example, workers that
|
||||
do not do the copying and only set their paths (only one worker per node does
|
||||
the copying) may try to execute functions that use the files in the directory
|
||||
being copied before the directory being copied has finished untarring.
|
||||
|
||||
Args:
|
||||
source_dir (str): The directory to copy.
|
||||
target_dir (str): The location to copy it to on the other machines. If this
|
||||
is not provided, the source_dir will be used. If it is provided and is
|
||||
different from source_dir, the source_dir also be copied to the target_dir
|
||||
location on this machine.
|
||||
"""
|
||||
target_dir = source_dir if target_dir is None else target_dir
|
||||
source_dir = os.path.abspath(source_dir)
|
||||
target_dir = os.path.abspath(target_dir)
|
||||
source_basename = os.path.basename(source_dir)
|
||||
target_basename = os.path.basename(target_dir)
|
||||
if source_basename != target_basename:
|
||||
raise Exception("The source_dir and target_dir must have the same base name, {} != {}".format(source_basename, target_basename))
|
||||
tarred_bytes = tarred_directory_as_bytes(source_dir)
|
||||
def f(worker_info):
|
||||
if worker_info["counter"] == 0:
|
||||
tarred_bytes_to_directory(tarred_bytes, os.path.dirname(target_dir))
|
||||
sys.path.append(os.path.dirname(target_dir))
|
||||
# Run this function on all workers to copy the directory to all nodes and to
|
||||
# add the directory to the Python path of each worker.
|
||||
ray.worker.global_worker.run_function_on_all_workers(f)
|
||||
Reference in New Issue
Block a user