Add files via upload

2024-01-16 17:38:48 +08:00
parent 143d32f621
commit 41ca6028d6
65 changed files with 139856 additions and 0 deletions
--- a/GPT_SoVITS/AR/models/t2s_lightning_module.py
+++ b/GPT_SoVITS/AR/models/t2s_lightning_module.py
@@ -0,0 +1,128 @@
+# modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/model/t2s_lightning_module.py
+import os,sys
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+from typing import Dict
+
+import torch
+from pytorch_lightning import LightningModule
+from AR.models.t2s_model import Text2SemanticDecoder
+from AR.modules.lr_schedulers import WarmupCosineLRSchedule
+from AR.modules.optim import ScaledAdam
+
+
+class Text2SemanticLightningModule(LightningModule):
+    def __init__(self, config, output_dir,is_train=True):
+        super().__init__()
+        self.config = config
+        self.top_k = 3
+        self.model = Text2SemanticDecoder(config=config, top_k=self.top_k)
+        pretrained_s1=config.get("pretrained_s1")
+        if(pretrained_s1 and is_train):
+            # print(self.load_state_dict(torch.load(pretrained_s1,map_location="cpu")["state_dict"]))
+            print(self.load_state_dict(torch.load(pretrained_s1,map_location="cpu")["weight"]))
+        if is_train:
+            self.automatic_optimization = False
+            self.save_hyperparameters()
+            self.eval_dir = output_dir / 'eval'
+            self.eval_dir.mkdir(parents=True, exist_ok=True)
+
+    def training_step(self, batch: Dict, batch_idx: int):
+
+        opt = self.optimizers()
+        scheduler = self.lr_schedulers()
+        loss, acc = self.model.forward(
+            batch['phoneme_ids'], batch['phoneme_ids_len'],
+            batch['semantic_ids'], batch['semantic_ids_len'],
+            batch['bert_feature'])
+        self.manual_backward(loss)
+        if batch_idx > 0 and batch_idx % 4 == 0:
+            opt.step()
+            opt.zero_grad()
+            scheduler.step()
+
+        self.log(
+            "total_loss",
+            loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True)
+        self.log(
+            "lr",
+            scheduler.get_last_lr()[0],
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True)
+        self.log(
+            f"top_{self.top_k}_acc",
+            acc,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True)
+
+    def validation_step(self, batch: Dict, batch_idx: int):return
+        # # get loss
+        # loss, acc = self.model.forward(
+        #     batch['phoneme_ids'], batch['phoneme_ids_len'],
+        #     batch['semantic_ids'], batch['semantic_ids_len'],
+        #     batch['bert_feature']
+        # )
+        #
+        # self.log(
+        #     "val_total_loss",
+        #     loss,
+        #     on_step=True,
+        #     on_epoch=True,
+        #     prog_bar=True,
+        #     sync_dist=True)
+        # self.log(
+        #     f"val_top_{self.top_k}_acc",
+        #     acc,
+        #     on_step=True,
+        #     on_epoch=True,
+        #     prog_bar=True,
+        #     sync_dist=True)
+        #
+        # # get infer output
+        # semantic_len = batch['semantic_ids'].size(1)
+        # prompt_len = min(int(semantic_len * 0.5), 150)
+        # prompt = batch['semantic_ids'][:, :prompt_len]
+        # pred_semantic = self.model.infer(batch['phoneme_ids'],
+        #                                  batch['phoneme_ids_len'], prompt,
+        #                                  batch['bert_feature']
+        #                                  )
+        # save_name = f'semantic_toks_{batch_idx}.pt'
+        # save_path = os.path.join(self.eval_dir, save_name)
+        # torch.save(pred_semantic.detach().cpu(), save_path)
+
+    def configure_optimizers(self):
+        model_parameters = self.model.parameters()
+        parameters_names = []
+        parameters_names.append([
+            name_param_pair[0]
+            for name_param_pair in self.model.named_parameters()
+        ])
+        lm_opt = ScaledAdam(
+            model_parameters,
+            lr=0.01,
+            betas=(0.9, 0.95),
+            clipping_scale=2.0,
+            parameters_names=parameters_names,
+            show_dominant_parameters=False,
+            clipping_update_period=1000, )
+
+        return {
+            "optimizer": lm_opt,
+            "lr_scheduler": {
+                "scheduler":
+                WarmupCosineLRSchedule(
+                    lm_opt,
+                    init_lr=self.config['optimizer']['lr_init'],
+                    peak_lr=self.config['optimizer']['lr'],
+                    end_lr=self.config['optimizer']['lr_end'],
+                    warmup_steps=self.config['optimizer']['warmup_steps'],
+                    total_steps=self.config['optimizer']['decay_steps'])
+            }
+        }