def fun(x, p, f, backend):
if "keops" in backend:
x = LazyTensor(x)
p = LazyTensor(p)
f = LazyTensor(f)
X = x * (-2 * math.pi * 1j * p * f).exp()
return X.sum(dim=0)
def squared_distances(x, y, use_keops=False):
if use_keops and keops_available:
if x.dim() == 2:
x_i = LazyTensor(x[:, None, :]) # (N,1,D)
y_j = LazyTensor(y[None, :, :]) # (1,M,D)
elif x.dim() == 3: # Batch computation
x_i = LazyTensor(x[:, :, None, :]) # (B,N,1,D)
y_j = LazyTensor(y[:, None, :, :]) # (B,1,M,D)
else:
print("x.shape : ", x.shape)
raise ValueError("Incorrect number of dimensions")
return ((x_i - y_j)**2).sum(-1)
else:
if x.dim() == 2:
D_xx = (x * x).sum(-1).unsqueeze(1) # (N,1)
D_xy = torch.matmul(x, y.permute(1, 0)) # (N,D) @ (D,M) = (N,M)
D_yy = (y * y).sum(-1).unsqueeze(0) # (1,M)
elif x.dim() == 3: # Batch computation
D_xx = (x * x).sum(-1).unsqueeze(2) # (B,N,1)
D_xy = torch.matmul(x, y.permute(0, 2,
1)) # (B,N,D) @ (B,D,M) = (B,N,M)
D_yy = (y * y).sum(-1).unsqueeze(1) # (B,1,M)
else:
print("x.shape : ", x.shape)
raise ValueError("Incorrect number of dimensions")
return D_xx - 2 * D_xy + D_yy
def kNN_graph(x, k, metric="euclidean"):
"""
Pykeops implementation of a k nearest neighbor graph
:param x: array containing the dataset
:param k: number of neartest neighbors
:param metric: Metric used for distances of x, must be "euclidean" or "cosine".
:return: array of shape (len(x), k) containing the indices of the k nearest neighbors of each datapoint
"""
x = torch.tensor(x).to("cuda").contiguous()
x_i = LazyTensor(x[:, None])
x_j = LazyTensor(x[None])
if metric == "euclidean":
dists = ((x_i - x_j)**2).sum(-1)
elif metric == "cosine":
scalar_prod = (x_i * x_j).sum(-1)
norm_x_i = (x_i**2).sum(-1).sqrt()
norm_x_j = (x_j**2).sum(-1).sqrt()
dists = 1 - scalar_prod / (norm_x_i * norm_x_j)
else:
raise NotImplementedError(f"Metric {metric} is not implemented.")
knn_idx = dists.argKmin(
K=k + 1, dim=0
)[:,
1:] # use k+1 neighbours and omit first, which is just the point itself
return knn_idx
def brute_force(self, x, y, k=5):
if use_cuda:
torch.cuda.synchronize()
x_LT = LazyTensor(x.unsqueeze(0))
y_LT = LazyTensor(y.unsqueeze(1))
D_ij = ((y_LT - x_LT)**2).sum(-1)
return D_ij.argKmin(K=k, axis=1)
def c2st(x_train,y_train,x_test,y_test,K=1):
###
# Input:
# x_train = tensor of shape [B, 1, IMG_DIM, IMG_DIM], training images
# y_train = tensor of shape [B, 1], training labels (here "fake" or "real")
# x_test = tensor of shape [B, 1, IMG_DIM, IMG_DIM], test images
# y_test = tensor of shape [B, 1], test labels
# k = number of neighbors to use in kNN classification
#
# Output:
# error = error of the kNN classifier
###
use_cuda = torch.cuda.is_available()
N = x_train.shape[3]
train_size = x_train.shape[0]
test_size = x_test.shape[0]
x_train = x_train.reshape(train_size,N*N)
x_test = x_test.reshape(test_size,N*N)
y_train = y_train.reshape(-1)
y_test = y_test.reshape(-1)
X_i = LazyTensor(x_test[:, None, :]) # test set
X_j = LazyTensor(x_train[None, :, :]) # train set
D_ij = ((X_i - X_j) ** 2).sum(-1) # Symbolic matrix of squared L2 distances
ind_knn = D_ij.argKmin(K, dim=1)
lab_knn = y_train[ind_knn]
y_knn, _ = lab_knn.mode() # Compute the most likely label
error = (y_knn != y_test).float().mean().item()
if error > 0.5:
error = 0.5
pykeops.clean_pykeops()
return error
def soft_distances(x, y, batch_x, batch_y, smoothness=0.01, atomtypes=None):
"""Computes a soft distance function to the atom centers of a protein.
Implements Eq. (1) of the paper in a fast and numerically stable way.
Args:
x (Tensor): (N,3) atom centers.
y (Tensor): (M,3) sampling locations.
batch_x (integer Tensor): (N,) batch vector for x, as in PyTorch_geometric.
batch_y (integer Tensor): (M,) batch vector for y, as in PyTorch_geometric.
smoothness (float, optional): atom radii if atom types are not provided. Defaults to .01.
atomtypes (integer Tensor, optional): (N,6) one-hot encoding of the atom chemical types. Defaults to None.
Returns:
Tensor: (M,) values of the soft distance function on the points `y`.
"""
# Build the (N, M, 1) symbolic matrix of squared distances:
x_i = LazyTensor(x[:, None, :]) # (N, 1, 3) atoms
y_j = LazyTensor(y[None, :, :]) # (1, M, 3) sampling points
D_ij = ((x_i - y_j)**2).sum(-1) # (N, M, 1) squared distances
# Use a block-diagonal sparsity mask to support heterogeneous batch processing:
D_ij.ranges = diagonal_ranges(batch_x, batch_y)
if atomtypes is not None:
# Turn the one-hot encoding "atomtypes" into a vector of diameters "smoothness_i":
# (N, 6) -> (N, 1, 1) (There are 6 atom types)
atomic_radii = torch.cuda.FloatTensor([170, 110, 152, 155, 180, 190],
device=x.device)
atomic_radii = atomic_radii / atomic_radii.min()
atomtype_radii = atomtypes * atomic_radii[
None, :] # n_atoms, n_atomtypes
# smoothness = atomtypes @ atomic_radii # (N, 6) @ (6,) = (N,)
smoothness = torch.sum(smoothness * atomtype_radii,
dim=1,
keepdim=False) # n_atoms, 1
smoothness_i = LazyTensor(smoothness[:, None, None])
# Compute an estimation of the mean smoothness in a neighborhood
# of each sampling point:
# density = (-D_ij.sqrt()).exp().sum(0).view(-1) # (M,) local density of atoms
# smooth = (smoothness_i * (-D_ij.sqrt()).exp()).sum(0).view(-1) # (M,)
# mean_smoothness = smooth / density # (M,)
# soft_dists = -mean_smoothness * (
# (-D_ij.sqrt() / smoothness_i).logsumexp(dim=0)
# ).view(-1)
mean_smoothness = (-D_ij.sqrt()).exp().sum(0)
mean_smoothness_j = LazyTensor(mean_smoothness[None, :, :])
mean_smoothness = (smoothness_i * (-D_ij.sqrt()).exp() /
mean_smoothness_j) # n_atoms, n_points, 1
mean_smoothness = mean_smoothness.sum(0).view(-1)
soft_dists = -mean_smoothness * (
(-D_ij.sqrt() / smoothness_i).logsumexp(dim=0)).view(-1)
else:
soft_dists = -smoothness * (
(-D_ij.sqrt() / smoothness).logsumexp(dim=0)).view(-1)
return soft_dists
def _KMeans(self, x: torch.tensor):
''' KMeans with Pykeops to do binning of original data.
Args:
x[np.array] = data
k_means[int] = number of bins to build
n_iter[int] = number iterations of KMeans loop
Returns:
labels[np.array] = class labels for each point in x
clusters[np.array] = coordinates for each centroid
'''
N, D = x.shape
clusters = torch.clone(x[:self.k_means, :]) # initialization of clusters
x_i = LazyTensor(x[:, None, :])
for i in range(self.n_iter):
clusters_j = LazyTensor(clusters[None, :, :])
D_ij = ((x_i - clusters_j) ** 2).sum(-1) # points-clusters kernel
labels = D_ij.argmin(axis=1).reshape(N) # Points -> Nearest cluster
Ncl = torch.bincount(labels) # Class weights
for d in range(D): # Compute the cluster centroids with np.bincount:
clusters[:, d] = torch.bincount(labels, weights=x[:, d]) / Ncl
return labels, clusters
def _pairwise_kernels(self, x: torch.tensor, y: torch.tensor = None, kernel='rbf',
sigma: float = 1.) -> LazyTensor:
'''Helper function to build kernel
Args: X = torch tensor of dimension 2.
K_type = type of Kernel to return
Returns:
K_ij[LazyTensor]
'''
if y is None:
y = x
if kernel == 'linear':
K_ij = x @ y.T
elif kernel == 'rbf':
x /= sigma
y /= sigma
x_i, x_j = LazyTensor(x[:, None, :]), LazyTensor(y[None, :, :])
K_ij = (-1 * ((x_i - x_j) ** 2).sum(-1)).exp()
# block-sparse reduction preprocess
K_ij = self._Gauss_block_sparse_pre(x, y, K_ij)
elif kernel == 'exp':
x_i, x_j = LazyTensor(x[:, None, :]), LazyTensor(y[None, :, :])
K_ij = (-1 * ((x_i - x_j) ** 2).sum().sqrt()).exp()
# block-sparse reduction preprocess
K_ij = self._Gauss_block_sparse_pre(x, y, K_ij) # TODO
K_ij = K_ij @ torch.diag(torch.ones(K_ij.shape[1])) # make 1 on diag only
K_ij.backend = self.backend
return K_ij
def _kNN(x, y, K):
"""Get K nearest neighbours using keops"""
x_i = LazyTensor(x[:, None, :]) # (M, 1, k)
y_j = LazyTensor(y[None, :, :]) # (1, N, k)
D_ij = ((x_i - y_j)**2).sum(
-1) # (M, N) symbolic matrix of squared distances
return D_ij.argKmin(K, dim=1) # (M, K) Minimum indices
def cluster(self, x, center=None):
if center is None:
center = self.InitFunc(x) ## M*D
with torch.no_grad():
x_lazy = LazyTensor(x.unsqueeze(1)) ## (N,1,D)
cl = None
for iter in range(self.iters):
c_lazy = LazyTensor(center.unsqueeze(0)) ## 1*M*D
dist = ((x_lazy - c_lazy)**2).sum(-1) ## N*M
cl = dist.argmin(dim=1).view(-1)
if iter < self.iters - 1:
for cnt_id in range(self.num_cnt):
selected = torch.nonzero(cl == cnt_id).squeeze(-1)
if selected.shape[0] != 0:
selected = torch.index_select(x, 0, selected)
center[cnt_id] = selected.mean(dim=0)
center_new = torch.zeros_like(center).to(torch.device('cuda'))
for cnt_id in range(self.num_cnt):
selected = torch.nonzero(cl == cnt_id).squeeze(-1)
if selected.shape[0] != 0:
selected = torch.index_select(x, 0, selected)
center_new[cnt_id] = selected.mean(dim=0)
return center_new
def fit(self, X:LazyTensor):
'''
Args: X = torch lazy tensor with features of shape
(1, n_samples, n_features)
Returns: Fitted instance of the class
'''
# Basic checks: we have a lazy tensor and n_components isn't too large
assert type(X) == LazyTensor, 'Input to fit(.) must be a LazyTensor.'
assert X.shape[1] >= self.n_components, f'The application needs X.shape[1] >= n_components.'
X = X.sum(dim=0)
# Number of samples
n_samples = X.size(0)
# Define basis
rnd = check_random_state(self.random_state)
inds = rnd.permutation(n_samples)
basis_inds = inds[:self.n_components]
basis = X[basis_inds]
# Build smaller kernel
basis_kernel = self._pairwise_kernels(basis, kernel = self.kernel)
# Get SVD
U, S, V = torch.svd(basis_kernel)
S = torch.maximum(S, torch.ones(S.size()) * 1e-12)
self.normalization_ = torch.mm(U / np.sqrt(S), V.t())
self.components_ = basis
self.component_indices_ = inds
return self
def KMeans(x, K=10, Niter=10, verbose=True):
N, D = x.shape # Number of samples, dimension of the ambient space
# K-means loop:
# - x is the point cloud,
# - cl is the vector of class labels
# - c is the cloud of cluster centroids
start = time.time()
c = x[:K, :].clone() # Simplistic random initialization
x_i = LazyTensor(x[:, None, :]) # (Npoints, 1, D)
for i in range(Niter):
c_j = LazyTensor(c[None, :, :]) # (1, Nclusters, D)
D_ij = ((x_i - c_j)**2).sum(
-1) # (Npoints, Nclusters) symbolic matrix of squared distances
cl = D_ij.argmin(dim=1).long().view(-1) # Points -> Nearest cluster
Ncl = torch.bincount(cl).type(torchtype[dtype]) # Class weights
for d in range(
D): # Compute the cluster centroids with torch.bincount:
c[:, d] = torch.bincount(cl, weights=x[:, d]) / Ncl
end = time.time()
if verbose:
print("K-means example with {:,} points in dimension {:,}, K = {:,}:".
format(N, D, K))
print('Timing for {} iterations: {:.5f}s = {} x {:.5f}s\n'.format(
Niter, end - start, Niter, (end - start) / Niter))
return cl, c
def local_moments(points, radius=1, ranges=None):
# print(radius)
# B, N, D = points.shape
shape_head, D = points.shape[:-1], points.shape[-1]
scale = 1.41421356237 * radius # math.sqrt(2) is not super JIT-friendly...
x = points / scale # Normalize the kernel size
# Computation:
x = torch.cat((torch.ones_like(x[..., :1]), x), dim=-1) # (B, N, D+1)
x_i = LazyTensor(x[..., :, None, :]) # (B, N, 1, D+1)
x_j = LazyTensor(x[..., None, :, :]) # (B, 1, N, D+1)
D_ij = ((x_i - x_j)**2).sum(-1) # (B, N, N), squared distances
K_ij = (-D_ij).exp() # (B, N, N), Gaussian kernel
C_ij = (K_ij * x_j).tensorprod(x_j) # (B, N, N, (D+1)*(D+1))
C_ij.ranges = ranges
C_i = C_ij.sum(dim=len(shape_head)).view(
shape_head +
(D + 1, D + 1)) # (B, N, D+1, D+1) : descriptors of order 0, 1 and 2
w_i = C_i[..., :1, :1] # (B, N, 1, 1), weights
m_i = C_i[..., :1, 1:] * scale # (B, N, 1, D), sum
c_i = C_i[..., 1:, 1:] * (scale**2) # (B, N, D, D), outer products
mass_i = w_i.squeeze(-1).squeeze(-1) # (B, N)
dev_i = (m_i / w_i).squeeze(-2) - points # (B, N, D)
cov_i = (c_i - (m_i.transpose(-1, -2) * m_i) / w_i) / w_i # (B, N, D, D)
return mass_i, dev_i, cov_i
def GaussKernelSum(x, y, gpuid=-1):
x_i = LazyTensor(x[:, None, :]) # x_i.shape = (M, 1, 3)
y_j = LazyTensor(y[None, :, :]) # y_j.shape = (1, N, 3)
D_ij = ((x_i - y_j)**2).sum(
dim=2) # Symbolic (M,N,1) matrix of squared distances
K_ij = (-D_ij).exp() # Symbolic (M,N,1) Gaussian kernel matrix
return K_ij.sum(dim=1, device_id=gpuid)
def predict(self, X):
r"""
Predict cluster belonging based on which cluster center is the closest.
Inputs:
:param X: torch.Tensor, (n_samples, n_features)
Outputs:
torch.Tensor, (n_samples,) int, indices identifying for each
datapoint which cluster center it belongs to.
"""
metric_function = self._get_distance_metric(self.distance_metric)
N_, _ = X.shape
K_, _ = self.cluster_centers_.shape
if K_ == 1:
# All points will belong to cluster 0.
return torch.zeros(N_)
if self.use_keops:
points_i = LazyTensor(X[None, :, None, :]) # (B=1, N, 1, C)
points_j = LazyTensor(
self.cluster_centers_[None, None, :, :]) # (B=1, 1, K, C)
cluster_id_ = metric_function(
points_i, points_j).argmin(dim=2)[0, :, 0] # rm B,C dims
else:
# TODO Untested
distances_to_centers = metric_function(
X[:, None, :], self.cluster_centers_[None, :, :]) # (N,K)
cluster_id_ = torch.argmin(distances_to_centers, dim=1)
return cluster_id_
def calc_centroid(x, c, cl, n=10):
"Helper function to optimise centroid location"
c = torch.clone(c.detach()).to(device)
c.requires_grad = True
x1 = LazyTensor(x.unsqueeze(0))
op = torch.optim.Adam([c], lr=1 / n)
scaling = 1 / torch.gather(torch.bincount(cl), 0, cl).view(-1, 1)
scaling.requires_grad = False
with torch.autograd.set_detect_anomaly(True):
for _ in range(n):
c.requires_grad = True
op.zero_grad()
c1 = LazyTensor(torch.index_select(c, 0, cl).unsqueeze(0))
d = distance(x1, c1)
loss = (d.sum(0) * scaling).sum(
) # calculate distance to centroid for each datapoint, divide by total number of points in that cluster, and sum
loss.backward(retain_graph=False)
op.step()
if normalise:
with torch.no_grad():
c = c / torch.norm(c, dim=-1).repeat_interleave(
c.shape[1]).reshape(
-1, c.shape[1]
) # normalising centroids to have norm 1
return c.detach()
def nlog_density(self, target, source, log_weights, scales, probas):
"""Negative log-likelihood of the proposal generated by the source onto the target."""
x_i = LazyTensor(target[:, None, :]) # (N,1,D)
y_j = LazyTensor(source[None, :, :]) # (1,M,D)
D_ij = ((x_i - y_j) ** 2).sum(-1)
D_ij = 1 + D_ij / (2 * x_i[self.D - 1] * y_j[self.D - 1])
D_ij = (D_ij + (D_ij ** 2 - 1).sqrt()).log() # (N,M)
neighbors_ij = (scales - D_ij).step() # 1 if |x_i-y_j| <= scales, 0 otherwise
if self.D == 2:
volumes = float(np.pi) * (scales.exp() + (-scales).exp())
else:
raise NotImplementedError()
neighbors_ij = neighbors_ij / volumes
if log_weights is None:
neighbors_ij = neighbors_ij / len(source)
else:
w_j = LazyTensor(log_weights[None, :, None].exp())
neighbors_ij = neighbors_ij * w_j
densities_i = neighbors_ij.sum(axis=1) # (N,K)
return -(densities_i * probas[None, :]).sum(dim=1).log().view(-1)
def forward(self,
query_layer,
key_layer,
value_layer,
attention_mask=None):
# Expect (..., t, d) shape for query, key, value
# Expect (..., t) for mask (default: (n, 1, 1, t))
d = query_layer.shape[-1]
t = query_layer.shape[-2]
sizes = query_layer.size()
reduction_dim = len(query_layer.size()) - 1
q = LazyTensor(query_layer.unsqueeze(-2).contiguous() / math.sqrt(d))
k = LazyTensor(key_layer.unsqueeze(-3).contiguous())
v = LazyTensor(value_layer.unsqueeze(-3).contiguous())
attention_scores = (q | k)
if attention_mask is not None:
mask = LazyTensor(attention_mask.unsqueeze(-1))
attention_scores = attention_scores + mask
return (attention_scores).sumsoftmaxweight(
v, dim=reduction_dim).reshape(*sizes)
def kneighbors(self, y):
'''
Obtain the k nearest neighbors of the query dataset y
'''
if self.__x is None:
raise ValueError(
'Input dataset not fitted yet! Call .fit() first!')
if type(y) != torch.Tensor:
raise ValueError("Query dataset must be a torch tensor")
if y.device != self.__device:
raise ValueError(
'Input dataset and query dataset must be on same device')
if len(y.shape) != 2:
raise ValueError('Query dataset must be a 2D tensor')
if self.__x.shape[-1] != y.shape[-1]:
raise ValueError('Query and dataset must have same dimensions')
if use_cuda:
torch.cuda.synchronize()
y = y.contiguous()
y_labels = self.__assign(y)
y_ranges, _, _ = cluster_ranges_centroids(y, y_labels)
self.__y_ranges = y_ranges
y, y_labels = self.__sort_clusters(y, y_labels, store_x=False)
x_LT = LazyTensor(self.__x.unsqueeze(0).to(self.__device).contiguous())
y_LT = LazyTensor(y.unsqueeze(1).to(self.__device).contiguous())
D_ij = ((y_LT - x_LT)**2).sum(-1)
ranges_ij = from_matrix(y_ranges, self.__x_ranges, self.__keep)
D_ij.ranges = ranges_ij
nn = D_ij.argKmin(K=self.__k, axis=1)
return self.__unsort(nn)
def gaussian_kernel(x, y, sigma=.1):
x_i = LazyTensor(x[:, None, :]) # (M, 1, 1)
y_j = LazyTensor(y[None, :, :]) # (1, N, 1)
D_ij = ((x_i - y_j)**2).sum(
-1) # (M, N) symbolic matrix of squared distances
return (-D_ij /
(2 * sigma**2)).exp() # (M, N) symbolic Gaussian kernel matrix
def laplacian_kernel(x, y, sigma=.1):
x_i = LazyTensor(x[:, None, :]) # (M, 1, 1)
y_j = LazyTensor(y[None, :, :]) # (1, N, 1)
D_ij = ((x_i - y_j)**2).sum(
-1) # (M, N) symbolic matrix of squared distances
return (-D_ij.sqrt() /
sigma).exp() # (M, N) symbolic Laplacian kernel matrix
def forward(self, mv_points, points0=None):
# set points0 to mv_points if no points0 are specified
points0 = mv_points if points0 is None else points0
# input dimensions are B(atch), E(mbd dim), ? (other dimensions)
# pykeops expects embedding/channel dimension at the back and only one dimension for the points of one instance
mv_points = mv_points.view(*(mv_points.size()[:2]), -1).permute(0, 2, 1).contiguous()
mv_points_i = LazyTensor(mv_points[:, :, None, :]) # B N 1 E
mv_points_j = LazyTensor(mv_points[:, None, :, :]) # B 1 N E
diff_mv_points_ij = mv_points_i - mv_points_j # B N N E
dist_mv_points_ij = (diff_mv_points_ij**2).sum(-1).sqrt() # B N N (1)
if points0 is not None:
points0 = points0.view(*(points0.size()[:2]), -1).permute(0, 2, 1).contiguous()
points0_i = LazyTensor(points0[:, :, None, :]) # B N 1 E
points0_j = LazyTensor(points0[:, None, :, :]) # B 1 N E
diff_points0_ij = points0_i - points0_j # B N N E
dist_points0_ij = (diff_points0_ij**2).sum(-1).sqrt() # B N N (1)
else:
dist_points0_ij = dist_mv_points_ij
rho_ij = self.rho(dist_mv_points_ij, max_d=self.config_rho["d"]) # B N N (1)
w_ij = self.w(dist_points0_ij, beta=self.config_w["beta"]) # B N N (1)
loss_i = 0.5 * ((2 * rho_ij - 1) * w_ij).sum(2) # B N (1) 1, last PyKeOps reduction step, output is torch.tensor
loss_i = ContiguousBackward().apply(loss_i) # backward pass needs to be contiguous as well
loss = loss_i.sum(1) # B (1, 1, 1) final reduction
loss = loss.squeeze(0).squeeze(0) # remove spurious dimension due to keops
return 0.5*loss # half the loss as all distances are considered twice
def keops_knn(
x: torch.Tensor,
y: torch.Tensor,
k: int,
batch_x: Optional[torch.Tensor] = None,
batch_y: Optional[torch.Tensor] = None,
cosine: bool = False,
) -> torch.Tensor:
r"""Straightforward modification of PyTorch_geometric's knn method."""
x = x.view(-1, 1) if x.dim() == 1 else x
y = y.view(-1, 1) if y.dim() == 1 else y
y_i = LazyTensor(y[:, None, :])
x_j = LazyTensor(x[None, :, :])
if cosine:
D_ij = -(y_i | x_j)
else:
D_ij = ((y_i - x_j)**2).sum(-1)
D_ij.ranges = diagonal_ranges(batch_y, batch_x)
idy = D_ij.argKmin(k, dim=1) # (N, K)
rows = torch.arange(k * len(y), device=idy.device) // k
return torch.stack([rows, idy.view(-1)], dim=0)
def kmeans(x,
distance=None,
K=10,
Niter=10,
device="cuda",
approx=False,
n=10):
from pykeops.torch import LazyTensor
if distance is None:
distance = torchtools.distance_function("euclidean")
def calc_centroid(x, c, cl, n=10):
"Helper function to optimise centroid location"
c = torch.clone(c.detach()).to(device)
c.requires_grad = True
x1 = LazyTensor(x.unsqueeze(0))
op = torch.optim.Adam([c], lr=1 / n)
scaling = 1 / torch.gather(torch.bincount(cl), 0, cl).view(-1, 1)
scaling.requires_grad = False
with torch.autograd.set_detect_anomaly(True):
for _ in range(n):
c.requires_grad = True
op.zero_grad()
c1 = LazyTensor(torch.index_select(c, 0, cl).unsqueeze(0))
d = distance(x1, c1)
loss = (d.sum(0) * scaling).sum(
) # calculate distance to centroid for each datapoint, divide by total number of points in that cluster, and sum
loss.backward(retain_graph=False)
op.step()
return c.detach()
N, D = x.shape
c = x[:K, :].clone()
x_i = LazyTensor(x.view(N, 1, D).to(device))
for i in range(Niter):
c_j = LazyTensor(c.view(1, K, D).to(device))
D_ij = distance(x_i, c_j)
cl = D_ij.argmin(dim=1).long().view(-1)
# updating c: either with approximation or exact
if approx:
# approximate with GD optimisation
c = calc_centroid(x, c, cl, n)
else:
# exact from average
c.zero_()
c.scatter_add_(0, cl[:, None].repeat(1, D), x)
Ncl = torch.bincount(cl, minlength=K).type_as(c).view(K, 1)
c /= Ncl
if torch.any(torch.isnan(c)):
raise ValueError(
"NaN detected in centroids during KMeans, please check metric is correct"
)
return cl, c
def mean_shift_step(points,
reference=None,
bandwidth=1,
kernel='gaussian',
use_keops=True):
"""
Perform one mean shift step: Move points towards maxima of Kernel density estimate using reference.
:param points: torch.FloatTensor or torch.cuda.FloatTensor
Points that will be shifted. Shape should be arbitrary barch dimensions (B) followed by number of points (N) and
:param reference: torch.FloatTensor
Points used for the kernel density estimate. Shape should be identical to points, except in (N) dimension.
:param bandwidth: float
Bandwidth of the kernel to be used, i.e. standard deviation for gaussian, radius for flat kernel.
:param kernel: str
Which kernel to use. Has to be one of MeanShiftStep.KERNELS .
:param use_keops: bool
Whether to use the PyKeOps Library or only vanilla PyTorch.
:return: torch.FloatTensor
Shifted points. Shape identical to points.
"""
reference = points if reference is None else reference
assert reference.shape[-1] == points.shape[
-1], f'{reference.shape}, {points.shape}'
points_i = points[:, :, None, :] # B N1 1 E
points_j = reference[:, None, :, :] # B 1 N2 E
if use_keops:
points_i = LazyTensor(points_i)
points_j = LazyTensor(points_j)
# array of vector differences
v_ij = points_i - points_j # B N1 N2 E
# array of squared distances
s_ij = (v_ij**2).sum(-1) # B N1 N2 (1)
if kernel == MeanShiftStep.GAUSSIAN_KERNEL:
factor = points.new([
-1 / (2 * bandwidth**2)
]) # uses the shape, device, dtype of points with new data
k_ij = (s_ij * factor).exp() # B N1 N2 (1)
elif kernel == MeanShiftStep.FLAT_KERNEL:
squared_radius = points.new([bandwidth**2])
if use_keops:
k_ij = (-s_ij + squared_radius).step() # B N1 N2 (1)
else:
k_ij = ((-s_ij + squared_radius) > 0).float() # B N1 N2 (1)
elif kernel == MeanShiftStep.EPANECHNIKOV_KERNEL:
squared_radius = points.new([bandwidth**2])
k_ij = (-s_ij + squared_radius).relu() # B N1 N2 (1)
else:
assert False, f'Kernel {kernel} not supported. Choose one of {MeanShiftStep.KERNELS}'
if not use_keops:
k_ij = k_ij.unsqueeze(-1) # KeOps never squeezes last dim
nominator = k_ij * points_j # B N1 N2 (1)
return nominator.sum(2) / k_ij.sum(2) # B N1 (1)
def K(x, y, b, **kwargs):
x_i = LazyTensor(x[:, None, :])
y_j = LazyTensor(y[None, :, :])
b_j = LazyTensor(b[None, :, :])
D_ij = (3.5 * (x_i - y_j)**2).sum(axis=2)
K_ij = (D_ij).sqrt() * b_j
K_ij = K_ij.sum(axis=1, call=False, **kwargs)
return K_ij
def gaussianconv_lazytensor(x, y, b, backend="GPU", **kwargs):
"""(B,N,D), (B,N,D), (B,N,1) -> (B,N,1)"""
x_i = LazyTensor(x.unsqueeze(-2)) # (B, M, 1, D)
y_j = LazyTensor(y.unsqueeze(-3)) # (B, 1, N, D)
D_ij = ((x_i - y_j) ** 2).sum(-1) # (B, M, N, 1)
K_ij = (-D_ij).exp() # (B, M, N, 1)
S_ij = K_ij * b.unsqueeze(-3) # (B, M, N, 1) * (B, 1, N, 1)
return S_ij.sum(dim=2, backend=backend)
def gaussianconv_lazytensor(x, y, b):
nbatchdims = len(x.shape) - 2
x_i = LazyTensor(x.unsqueeze(-2)) # (B, M, 1, D)
y_j = LazyTensor(y.unsqueeze(-3)) # (B, 1, N, D)
D_ij = ((x_i - y_j)**2).sum(-1) # (B, M, N, 1)
K_ij = (-D_ij).exp() # (B, M, N, 1)
S_ij = K_ij * b.unsqueeze(-3) # (B, M, N, 1) * (B, 1, N, 1)
return S_ij.sum(dim=nbatchdims + 1)
def forward(ctx, a, b):
one_hot = b.shape[-1]
a_lazy = LazyTensor(a.unsqueeze(-1).unsqueeze(-1).contiguous())
b_lazy = LazyTensor(b.unsqueeze(-1).unsqueeze(-1).contiguous())
c = (a_lazy + b_lazy).sum(-1).max(a.dim() - 1).squeeze(-1).squeeze(-1)
ac = (a_lazy + b_lazy).sum(-1).argmax(a.dim() - 1).squeeze(-1).squeeze(-1)
ctx.save_for_backward(ac, torch.tensor(one_hot))
return c
def nn_search(x_i, y_j, ranges=None):
x_i = LazyTensor(x_i[:, None, :]) # Broadcasted "line" variable
y_j = LazyTensor(y_j[None, :, :]) # Broadcasted "column" variable
D_ij = ((x_i - y_j)**2).sum(-1) # Symbolic matrix of squared distances
D_ij.ranges = ranges # Apply our block-sparsity pattern
return D_ij.argmin(dim=1).view(-1)
