Spaces:
Running
Running
import torch | |
import pytorch_lightning as pl | |
from torch.optim.lr_scheduler import ReduceLROnPlateau | |
from collections.abc import MutableMapping | |
from speechbrain.processing.speech_augmentation import SpeedPerturb | |
def flatten_dict(d, parent_key="", sep="_"): | |
"""Flattens a dictionary into a single-level dictionary while preserving | |
parent keys. Taken from | |
`SO <https://stackoverflow.com/questions/6027558/flatten-nested-dictionaries-compressing-keys>`_ | |
Args: | |
d (MutableMapping): Dictionary to be flattened. | |
parent_key (str): String to use as a prefix to all subsequent keys. | |
sep (str): String to use as a separator between two key levels. | |
Returns: | |
dict: Single-level dictionary, flattened. | |
""" | |
items = [] | |
for k, v in d.items(): | |
new_key = parent_key + sep + k if parent_key else k | |
if isinstance(v, MutableMapping): | |
items.extend(flatten_dict(v, new_key, sep=sep).items()) | |
else: | |
items.append((new_key, v)) | |
return dict(items) | |
class AudioLightningModule(pl.LightningModule): | |
def __init__( | |
self, | |
audio_model=None, | |
video_model=None, | |
optimizer=None, | |
loss_func=None, | |
train_loader=None, | |
val_loader=None, | |
test_loader=None, | |
scheduler=None, | |
config=None, | |
): | |
super().__init__() | |
self.audio_model = audio_model | |
self.video_model = video_model | |
self.optimizer = optimizer | |
self.loss_func = loss_func | |
self.train_loader = train_loader | |
self.val_loader = val_loader | |
self.test_loader = test_loader | |
self.scheduler = scheduler | |
self.config = {} if config is None else config | |
# Speed Aug | |
self.speedperturb = SpeedPerturb( | |
self.config["datamodule"]["data_config"]["sample_rate"], | |
speeds=[95, 100, 105], | |
perturb_prob=1.0 | |
) | |
# Save lightning"s AttributeDict under self.hparams | |
self.default_monitor = "val_loss/dataloader_idx_0" | |
self.save_hyperparameters(self.config_to_hparams(self.config)) | |
# self.print(self.audio_model) | |
self.validation_step_outputs = [] | |
self.test_step_outputs = [] | |
def forward(self, wav, mouth=None): | |
"""Applies forward pass of the model. | |
Returns: | |
:class:`torch.Tensor` | |
""" | |
return self.audio_model(wav) | |
def training_step(self, batch, batch_nb): | |
mixtures, targets, _ = batch | |
new_targets = [] | |
min_len = -1 | |
if self.config["training"]["SpeedAug"] == True: | |
with torch.no_grad(): | |
for i in range(targets.shape[1]): | |
new_target = self.speedperturb(targets[:, i, :]) | |
new_targets.append(new_target) | |
if i == 0: | |
min_len = new_target.shape[-1] | |
else: | |
if new_target.shape[-1] < min_len: | |
min_len = new_target.shape[-1] | |
targets = torch.zeros( | |
targets.shape[0], | |
targets.shape[1], | |
min_len, | |
device=targets.device, | |
dtype=torch.float, | |
) | |
for i, new_target in enumerate(new_targets): | |
targets[:, i, :] = new_targets[i][:, 0:min_len] | |
mixtures = targets.sum(1) | |
# print(mixtures.shape) | |
est_sources = self(mixtures) | |
loss = self.loss_func["train"](est_sources, targets) | |
self.log( | |
"train_loss", | |
loss, | |
on_epoch=True, | |
prog_bar=True, | |
sync_dist=True, | |
logger=True, | |
) | |
return {"loss": loss} | |
def validation_step(self, batch, batch_nb, dataloader_idx): | |
# cal val loss | |
if dataloader_idx == 0: | |
mixtures, targets, _ = batch | |
# print(mixtures.shape) | |
est_sources = self(mixtures) | |
loss = self.loss_func["val"](est_sources, targets) | |
self.log( | |
"val_loss", | |
loss, | |
on_epoch=True, | |
prog_bar=True, | |
sync_dist=True, | |
logger=True, | |
) | |
self.validation_step_outputs.append(loss) | |
return {"val_loss": loss} | |
# cal test loss | |
if (self.trainer.current_epoch) % 10 == 0 and dataloader_idx == 1: | |
mixtures, targets, _ = batch | |
# print(mixtures.shape) | |
est_sources = self(mixtures) | |
tloss = self.loss_func["val"](est_sources, targets) | |
self.log( | |
"test_loss", | |
tloss, | |
on_epoch=True, | |
prog_bar=True, | |
sync_dist=True, | |
logger=True, | |
) | |
self.test_step_outputs.append(tloss) | |
return {"test_loss": tloss} | |
def on_validation_epoch_end(self): | |
# val | |
avg_loss = torch.stack(self.validation_step_outputs).mean() | |
val_loss = torch.mean(self.all_gather(avg_loss)) | |
self.log( | |
"lr", | |
self.optimizer.param_groups[0]["lr"], | |
on_epoch=True, | |
prog_bar=True, | |
sync_dist=True, | |
) | |
self.logger.experiment.log( | |
{"learning_rate": self.optimizer.param_groups[0]["lr"], "epoch": self.current_epoch} | |
) | |
self.logger.experiment.log( | |
{"val_pit_sisnr": -val_loss, "epoch": self.current_epoch} | |
) | |
# test | |
if (self.trainer.current_epoch) % 10 == 0: | |
avg_loss = torch.stack(self.test_step_outputs).mean() | |
test_loss = torch.mean(self.all_gather(avg_loss)) | |
self.logger.experiment.log( | |
{"test_pit_sisnr": -test_loss, "epoch": self.current_epoch} | |
) | |
self.validation_step_outputs.clear() # free memory | |
self.test_step_outputs.clear() # free memory | |
def configure_optimizers(self): | |
"""Initialize optimizers, batch-wise and epoch-wise schedulers.""" | |
if self.scheduler is None: | |
return self.optimizer | |
if not isinstance(self.scheduler, (list, tuple)): | |
self.scheduler = [self.scheduler] # support multiple schedulers | |
epoch_schedulers = [] | |
for sched in self.scheduler: | |
if not isinstance(sched, dict): | |
if isinstance(sched, ReduceLROnPlateau): | |
sched = {"scheduler": sched, "monitor": self.default_monitor} | |
epoch_schedulers.append(sched) | |
else: | |
sched.setdefault("monitor", self.default_monitor) | |
sched.setdefault("frequency", 1) | |
# Backward compat | |
if sched["interval"] == "batch": | |
sched["interval"] = "step" | |
assert sched["interval"] in [ | |
"epoch", | |
"step", | |
], "Scheduler interval should be either step or epoch" | |
epoch_schedulers.append(sched) | |
return [self.optimizer], epoch_schedulers | |
# def lr_scheduler_step(self, scheduler, optimizer_idx, metric): | |
# if metric is None: | |
# scheduler.step() | |
# else: | |
# scheduler.step(metric) | |
def train_dataloader(self): | |
"""Training dataloader""" | |
return self.train_loader | |
def val_dataloader(self): | |
"""Validation dataloader""" | |
return [self.val_loader, self.test_loader] | |
def on_save_checkpoint(self, checkpoint): | |
"""Overwrite if you want to save more things in the checkpoint.""" | |
checkpoint["training_config"] = self.config | |
return checkpoint | |
def config_to_hparams(dic): | |
"""Sanitizes the config dict to be handled correctly by torch | |
SummaryWriter. It flatten the config dict, converts ``None`` to | |
``"None"`` and any list and tuple into torch.Tensors. | |
Args: | |
dic (dict): Dictionary to be transformed. | |
Returns: | |
dict: Transformed dictionary. | |
""" | |
dic = flatten_dict(dic) | |
for k, v in dic.items(): | |
if v is None: | |
dic[k] = str(v) | |
elif isinstance(v, (list, tuple)): | |
dic[k] = torch.tensor(v) | |
return dic | |