def ConjugateGradientSolver(binding, linop, b, eps=1e-6): # Conjugate gradient algorithm to solve linear system of the form # Ma=b where linop is a linear operation corresponding # to a symmetric and positive definite matrix tools = get_tools(binding) delta = tools.size(b) * eps**2 a = 0 r = tools.copy(b) nr2 = (r**2).sum() if nr2 < delta: return 0 * r p = tools.copy(r) k = 0 while True: Mp = linop(p) alp = nr2 / (p * Mp).sum() a += alp * p r -= alp * Mp nr2new = (r**2).sum() if nr2new < delta: break p = r + (nr2new / nr2) * p nr2 = nr2new k += 1 return a
def postprocess(out, binding, reduction_op, nout, opt_arg, dtype): tools = get_tools(binding) # Post-processing of the output: if "Arg" in reduction_op: # when using Arg type reductions, # if nout is greater than 16 millions and dtype=float32, the result is not reliable # because we encode indices as floats, so we raise an exception if nout > 1.6e7 and dtype == "float32": raise ValueError( 'size of input array is too large for Arg type reduction with single precision. Use double precision.' ) if reduction_op == 'SumSoftMaxWeight' or reduction_op == 'SoftMax': # we compute sum_j exp(f_ij) g_ij / sum_j exp(f_ij) from sum_j exp(m_i-f_ij) [1,g_ij] out = out[..., 2:] / out[..., 1][..., None] elif reduction_op == 'ArgMin' or reduction_op == 'ArgMax': # outputs are encoded as floats but correspond to indices, so we cast to integers out = tools.long(out) elif reduction_op == 'Min_ArgMin' or reduction_op == 'MinArgMin' or reduction_op == 'Max_ArgMax' or reduction_op == 'MaxArgMax': # output is one array of size N x 2D, giving min and argmin value for each dimension. # We convert to one array of floats of size NxD giving mins, and one array of size NxD giving argmins (casted to integers) shape_out = out.shape tmp = tools.view(out, shape_out[:-1] + (2, -1)) vals = tmp[..., 0, :] indices = tmp[..., 1, :] out = (vals, tools.long(indices)) elif reduction_op == 'KMin': # output is of size N x KD giving K minimal values for each dim. We convert to array of size N x K x D shape_out = out.shape out = tools.view(out, shape_out[:-1] + (opt_arg, -1)) if out.shape[-1] == 1: out = out.squeeze(-1) elif reduction_op == 'ArgKMin': # output is of size N x KD giving K minimal values for each dim. We convert to array of size N x K x D # and cast to integers shape_out = out.shape out = tools.view(tools.long(out), shape_out[:-1] + (opt_arg, -1)) if out.shape[-1] == 1: out = out.squeeze(-1) elif reduction_op == 'KMin_ArgKMin' or reduction_op == 'KMinArgKMin': # output is of size N x 2KD giving K min and argmin for each dim. We convert to 2 arrays of size N x K x D # and cast to integers the second array shape_out = out.shape out = tools.view(out, shape_out[:-1] + (opt_arg, 2, -1)) out = (out[..., 0, :], tools.long(out[..., 1, :])) if out[0].shape[-1] == 1: out = (out[0].squeeze(-1), out[1].squeeze(-1)) elif reduction_op == 'LogSumExp': # finalize the log-sum-exp computation as m + log(s) if out.shape[-1] == 2: # means (m,s) with m scalar and s scalar out = (out[..., 0] + tools.log(out[..., 1]))[..., None] else: # here out.shape[-1]>2, means (m,s) with m scalar and s vectorial out = out[..., 0][..., None] + tools.log(out[..., 1:]) return out
def postprocess(out, binding, reduction_op, nout, opt_arg, dtype): tools = get_tools(binding) # Post-processing of the output: if reduction_op == "SumSoftMaxWeight" or reduction_op == "SoftMax": # we compute sum_j exp(f_ij) g_ij / sum_j exp(f_ij) from sum_j exp(m_i-f_ij) [1,g_ij] out = out[..., 2:] / out[..., 1][..., None] elif reduction_op == "ArgMin" or reduction_op == "ArgMax": # outputs are encoded as floats but correspond to indices, so we cast to integers out = tools.long(out) elif (reduction_op == "Min_ArgMin" or reduction_op == "MinArgMin" or reduction_op == "Max_ArgMax" or reduction_op == "MaxArgMax"): # output is one array of size N x 2D, giving min and argmin value for each dimension. # We convert to one array of floats of size NxD giving mins, and one array of size NxD giving argmins (casted to integers) shape_out = out.shape tmp = tools.view(out, shape_out[:-1] + (2, -1)) vals = tmp[..., 0, :] indices = tmp[..., 1, :] out = (vals, tools.long(indices)) elif reduction_op == "KMin": # output is of size N x KD giving K minimal values for each dim. We convert to array of size N x K x D shape_out = out.shape out = tools.view(out, shape_out[:-1] + (opt_arg, -1)) if out.shape[-1] == 1: out = out.squeeze(-1) elif reduction_op == "ArgKMin": # output is of size N x KD giving K minimal values for each dim. We convert to array of size N x K x D # and cast to integers shape_out = out.shape out = tools.view(tools.long(out), shape_out[:-1] + (opt_arg, -1)) if out.shape[-1] == 1: out = out.squeeze(-1) elif reduction_op == "KMin_ArgKMin" or reduction_op == "KMinArgKMin": # output is of size N x 2KD giving K min and argmin for each dim. We convert to 2 arrays of size N x K x D # and cast to integers the second array shape_out = out.shape out = tools.view(out, shape_out[:-1] + (opt_arg, 2, -1)) out = (out[..., 0, :], tools.long(out[..., 1, :])) if out[0].shape[-1] == 1: out = (out[0].squeeze(-1), out[1].squeeze(-1)) elif reduction_op == "LogSumExp": # finalize the log-sum-exp computation as m + log(s) if out.shape[-1] == 2: # means (m,s) with m scalar and s scalar out = (out[..., 0] + tools.log(out[..., 1]))[..., None] else: # here out.shape[-1]>2, means (m,s) with m scalar and s vectorial out = out[..., 0][..., None] + tools.log(out[..., 1:]) return out
def KernelLinearSolver(binding, K, x, b, alpha=0, eps=1e-6, precond=False, precondKernel=None): tools = get_tools(binding) dtype = tools.dtype(x) def PreconditionedConjugateGradientSolver(linop, b, invprecondop, eps=1e-6): # Preconditioned conjugate gradient algorithm to solve linear system of the form # Ma=b where linop is a linear operation corresponding # to a symmetric and positive definite matrix # invprecondop is linear operation corresponding to the inverse of the preconditioner matrix a = 0 r = tools.copy(b) z = invprecondop(r) p = tools.copy(z) rz = (r * z).sum() k = 0 while True: alp = rz / (p * linop(p)).sum() a += alp * p r -= alp * linop(p) if (r**2).sum() < eps**2: break z = invprecondop(r) rznew = (r * z).sum() p = z + (rznew / rz) * p rz = rznew k += 1 return a def NystromInversePreconditioner(K, Kspec, x, alpha): N, D = x.shape m = int(np.sqrt(N)) ind = np.random.choice(range(N), m, replace=False) u = x[ind, :] M = K(u, u) + Kspec(tools.tile(u, (m, 1)), tools.tile(u, (1, m)).reshape(-1, D), x).reshape( m, m) def invprecondop(r): a = tools.solve(M, K(u, x, r)) return (r - K(x, u, a)) / alpha return invprecondop def KernelLinOp(a): return K(x, x, a) + alpha * a def GaussKernel(D, Dv, sigma): formula = "Exp(-oos2*SqDist(x,y))*b" variables = [ "x = Vi(" + str(D) + ")", # First arg : i-variable, of size D "y = Vj(" + str(D) + ")", # Second arg : j-variable, of size D "b = Vj(" + str(Dv) + ")", # Third arg : j-variable, of size Dv "oos2 = Pm(1)", ] # Fourth arg : scalar parameter my_routine = tools.Genred(formula, variables, reduction_op="Sum", axis=1, dtype=dtype) oos2 = tools.array([1.0 / sigma**2], dtype=dtype) KernelMatrix = GaussKernelMatrix(sigma) def K(x, y, b=None): if b is None: return KernelMatrix(x, y) else: return my_routine(x, y, b, oos2) return K def GaussKernelNystromPrecond(D, sigma): formula = "Exp(-oos2*(SqDist(u,x)+SqDist(v,x)))" variables = [ "u = Vi(" + str(D) + ")", # First arg : i-variable, of size D "v = Vi(" + str(D) + ")", # Second arg : i-variable, of size D "x = Vj(" + str(D) + ")", # Third arg : j-variable, of size D "oos2 = Pm(1)", ] # Fourth arg : scalar parameter my_routine = tools.Genred(formula, variables, reduction_op="Sum", axis=1, dtype=dtype) oos2 = tools.array([1.0 / sigma**2], dtype=dtype) KernelMatrix = GaussKernelMatrix(sigma) def K(u, v, x): return my_routine(u, v, x, oos2) return K def GaussKernelMatrix(sigma): oos2 = 1.0 / sigma**2 def f(x, y): D = x.shape[1] sqdist = 0 for k in range(D): sqdist += (x[:, k][:, None] - tools.transpose(y[:, k][:, None]))**2 return tools.exp(-oos2 * sqdist) return f if type(K) == tuple: if K[0] == "gaussian": D = K[1] Dv = K[2] sigma = K[3] K = GaussKernel(D, Dv, sigma) if precond: precondKernel = GaussKernelNystromPrecond(D, sigma) if precond: invprecondop = NystromInversePreconditioner(K, precondKernel, x, alpha) a = PreconditionedConjugateGradientSolver(KernelLinOp, b, invprecondop, eps) else: a = ConjugateGradientSolver(binding, KernelLinOp, b, eps=eps) return a