Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion viscy/representation/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
from lightning.pytorch import LightningModule
from pytorch_metric_learning.losses import NTXentLoss
from torch import Tensor, nn
from umap import UMAP

from viscy.data.typing import TrackingIndex, TripletSample
from viscy.representation.contrastive import ContrastiveEncoder
Expand Down Expand Up @@ -183,6 +182,8 @@ def _log_step_samples(self, batch_idx, samples, stage: Literal["train", "val"]):
output_list.extend(detach_sample(samples, self.log_samples_per_batch))

def log_embedding_umap(self, embeddings: Tensor, tag: str):
from umap import UMAP

_logger.debug(f"Computing UMAP for {tag} embeddings.")
umap = UMAP(n_components=2)
embeddings_np = embeddings.detach().cpu().numpy()
Expand Down
Loading