chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,211 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
from collections import Counter
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class EM:
|
||||
"""
|
||||
EM algorithm used to quantize the columns of W to minimize
|
||||
|
||||
||W - W_hat||^2
|
||||
|
||||
Args:
|
||||
- W: weight matrix of size (in_features x out_features)
|
||||
- n_iter: number of k-means iterations
|
||||
- n_centroids: number of centroids (size of codebook)
|
||||
- eps: for cluster reassignment when an empty cluster is found
|
||||
- max_tentatives for cluster reassignment when an empty cluster is found
|
||||
- verbose: print error after each iteration
|
||||
|
||||
Remarks:
|
||||
- If one cluster is empty, the most populated cluster is split into
|
||||
two clusters
|
||||
- All the relevant dimensions are specified in the code
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, W, n_centroids=256, n_iter=20, eps=1e-6, max_tentatives=30, verbose=True
|
||||
):
|
||||
self.W = W
|
||||
self.n_centroids = n_centroids
|
||||
self.n_iter = n_iter
|
||||
self.eps = eps
|
||||
self.max_tentatives = max_tentatives
|
||||
self.verbose = verbose
|
||||
self.centroids = torch.Tensor()
|
||||
self.assignments = torch.Tensor()
|
||||
self.objective = []
|
||||
|
||||
def initialize_centroids(self):
|
||||
"""
|
||||
Initializes the centroids by sampling random columns from W.
|
||||
"""
|
||||
|
||||
in_features, out_features = self.W.size()
|
||||
indices = torch.randint(
|
||||
low=0, high=out_features, size=(self.n_centroids,)
|
||||
).long()
|
||||
self.centroids = self.W[:, indices].t() # (n_centroids x in_features)
|
||||
|
||||
def step(self, i):
|
||||
"""
|
||||
There are two standard steps for each iteration: expectation (E) and
|
||||
minimization (M). The E-step (assignment) is performed with an exhaustive
|
||||
search and the M-step (centroid computation) is performed with
|
||||
the exact solution.
|
||||
|
||||
Args:
|
||||
- i: step number
|
||||
|
||||
Remarks:
|
||||
- The E-step heavily uses PyTorch broadcasting to speed up computations
|
||||
and reduce the memory overhead
|
||||
"""
|
||||
|
||||
# assignments (E-step)
|
||||
distances = self.compute_distances() # (n_centroids x out_features)
|
||||
self.assignments = torch.argmin(distances, dim=0) # (out_features)
|
||||
n_empty_clusters = self.resolve_empty_clusters()
|
||||
|
||||
# centroids (M-step)
|
||||
for k in range(self.n_centroids):
|
||||
W_k = self.W[:, self.assignments == k] # (in_features x size_of_cluster_k)
|
||||
self.centroids[k] = W_k.mean(dim=1) # (in_features)
|
||||
|
||||
# book-keeping
|
||||
obj = (self.centroids[self.assignments].t() - self.W).norm(p=2).item()
|
||||
self.objective.append(obj)
|
||||
if self.verbose:
|
||||
logging.info(
|
||||
f"Iteration: {i},\t"
|
||||
f"objective: {obj:.6f},\t"
|
||||
f"resolved empty clusters: {n_empty_clusters}"
|
||||
)
|
||||
|
||||
def resolve_empty_clusters(self):
|
||||
"""
|
||||
If one cluster is empty, the most populated cluster is split into
|
||||
two clusters by shifting the respective centroids. This is done
|
||||
iteratively for a fixed number of tentatives.
|
||||
"""
|
||||
|
||||
# empty clusters
|
||||
counts = Counter(map(lambda x: x.item(), self.assignments))
|
||||
empty_clusters = set(range(self.n_centroids)) - set(counts.keys())
|
||||
n_empty_clusters = len(empty_clusters)
|
||||
|
||||
tentatives = 0
|
||||
while len(empty_clusters) > 0:
|
||||
# given an empty cluster, find most populated cluster and split it into two
|
||||
k = random.choice(list(empty_clusters))
|
||||
m = counts.most_common(1)[0][0]
|
||||
e = torch.randn_like(self.centroids[m]) * self.eps
|
||||
self.centroids[k] = self.centroids[m].clone()
|
||||
self.centroids[k] += e
|
||||
self.centroids[m] -= e
|
||||
|
||||
# recompute assignments
|
||||
distances = self.compute_distances() # (n_centroids x out_features)
|
||||
self.assignments = torch.argmin(distances, dim=0) # (out_features)
|
||||
|
||||
# check for empty clusters
|
||||
counts = Counter(map(lambda x: x.item(), self.assignments))
|
||||
empty_clusters = set(range(self.n_centroids)) - set(counts.keys())
|
||||
|
||||
# increment tentatives
|
||||
if tentatives == self.max_tentatives:
|
||||
logging.info(
|
||||
f"Could not resolve all empty clusters, {len(empty_clusters)} remaining"
|
||||
)
|
||||
raise EmptyClusterResolveError
|
||||
tentatives += 1
|
||||
|
||||
return n_empty_clusters
|
||||
|
||||
def compute_distances(self):
|
||||
"""
|
||||
For every centroid m, computes
|
||||
|
||||
||M - m[None, :]||_2
|
||||
|
||||
Remarks:
|
||||
- We rely on PyTorch's broadcasting to speed up computations
|
||||
and reduce the memory overhead
|
||||
- Without chunking, the sizes in the broadcasting are modified as:
|
||||
(n_centroids x n_samples x out_features) -> (n_centroids x out_features)
|
||||
- The broadcasting computation is automatically chunked so that
|
||||
the tensors fit into the memory of the GPU
|
||||
"""
|
||||
|
||||
nb_centroids_chunks = 1
|
||||
|
||||
while True:
|
||||
try:
|
||||
return torch.cat(
|
||||
[
|
||||
(self.W[None, :, :] - centroids_c[:, :, None]).norm(p=2, dim=1)
|
||||
for centroids_c in self.centroids.chunk(
|
||||
nb_centroids_chunks, dim=0
|
||||
)
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
except RuntimeError:
|
||||
nb_centroids_chunks *= 2
|
||||
|
||||
def assign(self):
|
||||
"""
|
||||
Assigns each column of W to its closest centroid, thus essentially
|
||||
performing the E-step in train().
|
||||
|
||||
Remarks:
|
||||
- The function must be called after train() or after loading
|
||||
centroids using self.load(), otherwise it will return empty tensors
|
||||
"""
|
||||
|
||||
distances = self.compute_distances() # (n_centroids x out_features)
|
||||
self.assignments = torch.argmin(distances, dim=0) # (out_features)
|
||||
|
||||
def save(self, path, layer):
|
||||
"""
|
||||
Saves centroids and assignments.
|
||||
|
||||
Args:
|
||||
- path: folder used to save centroids and assignments
|
||||
"""
|
||||
|
||||
torch.save(self.centroids, os.path.join(path, "{}_centroids.pth".format(layer)))
|
||||
torch.save(
|
||||
self.assignments, os.path.join(path, "{}_assignments.pth".format(layer))
|
||||
)
|
||||
torch.save(self.objective, os.path.join(path, "{}_objective.pth".format(layer)))
|
||||
|
||||
def load(self, path, layer):
|
||||
"""
|
||||
Loads centroids and assignments from a given path
|
||||
|
||||
Args:
|
||||
- path: folder use to load centroids and assignments
|
||||
"""
|
||||
|
||||
self.centroids = torch.load(
|
||||
os.path.join(path, "{}_centroids.pth".format(layer))
|
||||
)
|
||||
self.assignments = torch.load(
|
||||
os.path.join(path, "{}_assignments.pth".format(layer))
|
||||
)
|
||||
self.objective = torch.load(
|
||||
os.path.join(path, "{}_objective.pth".format(layer))
|
||||
)
|
||||
|
||||
|
||||
class EmptyClusterResolveError(Exception):
|
||||
pass
|
||||
Reference in New Issue
Block a user