Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support MLflow along with Tensorboard for logging segmentation task visualizations during validation #1209

Merged
merged 1 commit into from Jan 3, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
16 changes: 12 additions & 4 deletions pyannote/audio/tasks/segmentation/mixins.py
Expand Up @@ -30,6 +30,7 @@
import numpy as np
import torch
from pyannote.core import Segment, SlidingWindowFeature
from pytorch_lightning.loggers import TensorBoardLogger, MLFlowLogger
from torch.utils.data._utils.collate import default_collate
from torchmetrics import Metric
from torchmetrics.classification import BinaryAUROC, MultilabelAUROC, MulticlassAUROC
Expand Down Expand Up @@ -429,7 +430,7 @@ def validation_step(self, batch, batch_idx: int):
):
return

# visualize first 9 validation samples of first batch in Tensorboard
# visualize first 9 validation samples of first batch in Tensorboard/MLflow
X = X.cpu().numpy()
y = y.float().cpu().numpy()
y_pred = y_pred.cpu().numpy()
Expand Down Expand Up @@ -478,8 +479,15 @@ def validation_step(self, batch, batch_idx: int):

plt.tight_layout()

self.model.logger.experiment.add_figure(
f"{self.logging_prefix}ValSamples", fig, self.model.current_epoch
)
if isinstance(self.model.logger, TensorBoardLogger):
self.model.logger.experiment.add_figure(
f"{self.logging_prefix}ValSamples", fig, self.model.current_epoch
)
elif isinstance(self.model.logger, MLFlowLogger):
self.model.logger.experiment.log_figure(
run_id=self.model.logger.run_id,
figure=fig,
artifact_file=f"{self.logging_prefix}ValSamples_epoch{self.model.current_epoch}.png",
)

plt.close(fig)