| | |
| |
|
| | import pytorch_lightning as pl |
| | import torch |
| | from sklearn.cluster import KMeans |
| | import numpy as np |
| |
|
| | class RBFNetwork(pl.LightningModule): |
| | def __init__( |
| | self, |
| | current_timestep, |
| | next_timestep, |
| | n_centers: int = 100, |
| | kappa: float = 1.0, |
| | lr=1e-2, |
| | datamodule=None, |
| | image_data=False, |
| | args=None |
| | ): |
| | super().__init__() |
| | self.K = n_centers |
| | self.current_timestep = current_timestep |
| | self.next_timestep = next_timestep |
| | self.clustering_model = KMeans(n_clusters=self.K) |
| | self.kappa = kappa |
| | self.last_val_loss = 1 |
| | self.lr = lr |
| | self.W = torch.nn.Parameter(torch.rand(self.K, 1)) |
| | self.datamodule = datamodule |
| | self.image_data = image_data |
| | self.args = args |
| |
|
| | def on_before_zero_grad(self, *args, **kwargs): |
| | self.W.data = torch.clamp(self.W.data, min=0.0001) |
| |
|
| | def on_train_start(self): |
| | with torch.no_grad(): |
| | |
| | batch = next(iter(self.trainer.datamodule.train_dataloader())) |
| | |
| | metric_samples = batch[0]["metric_samples"][0] |
| | all_data = torch.cat(metric_samples) |
| | data_to_fit = all_data |
| |
|
| | print("Fitting Clustering model...") |
| | self.clustering_model.fit(data_to_fit) |
| |
|
| | clusters = ( |
| | self.calculate_centroids(all_data, self.clustering_model.labels_) |
| | if self.image_data |
| | else self.clustering_model.cluster_centers_ |
| | ) |
| |
|
| | self.C = torch.tensor(clusters, dtype=torch.float32).to(self.device) |
| | labels = self.clustering_model.labels_ |
| | sigmas = np.zeros((self.K, 1)) |
| |
|
| | for k in range(self.K): |
| | points = all_data[labels == k, :] |
| | variance = ((points - clusters[k]) ** 2).mean(axis=0) |
| | sigmas[k, :] = np.sqrt( |
| | variance.sum() if self.image_data else variance.mean() |
| | ) |
| |
|
| | |
| | sigmas = np.maximum(sigmas, 1e-6) |
| | |
| | self.lamda = torch.tensor( |
| | 0.5 / (self.kappa * sigmas) ** 2, dtype=torch.float32 |
| | ).to(self.device) |
| |
|
| | def forward(self, x): |
| | if len(x.shape) > 2: |
| | x = x.reshape(x.shape[0], -1).to(self.C.device) |
| | |
| | x = x.to(self.C.device) |
| | dist2 = torch.cdist(x, self.C) ** 2 |
| | self.phi_x = torch.exp(-0.5 * self.lamda[None, :, :] * dist2[:, :, None]) |
| | |
| | h_x = (self.W.to(x.device) * self.phi_x).sum(dim=1) |
| | |
| | return h_x |
| |
|
| | def training_step(self, batch, batch_idx): |
| | if self.args.data_type == "scrna" or self.args.data_type == "tahoe": |
| | main_batch = batch[0]["train_samples"][0] |
| | else: |
| | main_batch = batch["train_samples"][0] |
| |
|
| | x0 = main_batch["x0"][0] |
| | if self.args.branches == 1: |
| | x1 = main_batch["x1"][0] |
| | inputs = torch.cat([x0, x1], dim=0).to(self.device) |
| | else: |
| | x1_1 = main_batch["x1_1"][0] |
| | x1_2 = main_batch["x1_2"][0] |
| | |
| | inputs = torch.cat([x0, x1_1, x1_2], dim=0).to(self.device) |
| | print("inputs shape") |
| | print(inputs.shape) |
| | |
| | loss = ((1 - self.forward(inputs)) ** 2).mean() |
| | self.log( |
| | "MetricModel/train_loss_learn_metric", |
| | loss, |
| | on_step=True, |
| | on_epoch=True, |
| | prog_bar=True, |
| | ) |
| | return loss |
| |
|
| | def validation_step(self, batch, batch_idx): |
| | if self.args.data_type == "scrna" or self.args.data_type == "tahoe": |
| | main_batch = batch[0]["val_samples"][0] |
| | else: |
| | main_batch = batch["val_samples"][0] |
| |
|
| | x0 = main_batch["x0"][0] |
| | if self.args.branches == 1: |
| | x1 = main_batch["x1"][0] |
| | inputs = torch.cat([x0, x1], dim=0).to(self.device) |
| | else: |
| | x1_1 = main_batch["x1_1"][0] |
| | x1_2 = main_batch["x1_2"][0] |
| | |
| | inputs = torch.cat([x0, x1_1, x1_2], dim=0).to(self.device) |
| |
|
| | h = self.forward(inputs) |
| | |
| | loss = ((1 - h) ** 2).mean() |
| | self.log( |
| | "MetricModel/val_loss_learn_metric", |
| | loss, |
| | on_step=True, |
| | on_epoch=True, |
| | prog_bar=True, |
| | ) |
| | self.last_val_loss = loss.detach() |
| | return loss |
| |
|
| | def calculate_centroids(self, all_data, labels): |
| | unique_labels = np.unique(labels) |
| | centroids = np.zeros((len(unique_labels), all_data.shape[1])) |
| | for i, label in enumerate(unique_labels): |
| | centroids[i] = all_data[labels == label].mean(axis=0) |
| | return centroids |
| |
|
| | def configure_optimizers(self): |
| | optimizer = torch.optim.Adam(self.parameters(), lr=self.lr) |
| | return optimizer |
| |
|
| | def compute_metric(self, x, alpha=1, epsilon=1e-2, image_hx=False): |
| | if epsilon < 0: |
| | epsilon = (1 - float(self.last_val_loss)) / abs(epsilon) |
| | h_x = self.forward(x) |
| | if image_hx: |
| | h_x = 1 - torch.abs(1 - h_x) |
| | M_x = 1 / (h_x**alpha + epsilon) |
| | else: |
| | M_x = 1 / (h_x + epsilon) ** alpha |
| | return M_x |
