File size: 3,573 Bytes
9e14838 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | import importlib
import json
import os
import warnings
warnings.filterwarnings("ignore")
from typing import Any, Dict, List, Optional, Tuple
import argparse
import pytorch_lightning as pl
import torch
import hydra
torch.set_float32_matmul_precision("high")
from pytorch_lightning import Callback, LightningDataModule, LightningModule, Trainer
from pytorch_lightning.strategies.ddp import DDPStrategy
from omegaconf import DictConfig
from omegaconf import OmegaConf
from pytorch_lightning.utilities import rank_zero_only
@rank_zero_only
def print_only(message: str):
"""Prints a message only on rank 0."""
print(message)
def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
# instantiate datamodule
print_only(f"Instantiating datamodule <{cfg.datamodule._target_}>")
datamodule: LightningDataModule = hydra.utils.instantiate(cfg.datamodule)
datamodule.setup()
# instantiate decouple model
print_only(f"Instantiating decouple model <{cfg.decouple_model._target_}>")
decouple_model: torch.nn.Module = hydra.utils.instantiate(cfg.decouple_model)
decouple_model.load_state_dict(torch.load(cfg.speechtokenizer_path))
# import pdb; pdb.set_trace()
# instantiate detect model
print(f"Instantiating detect model <{cfg.detect_model._target_}>")
detect_model: torch.nn.Module = hydra.utils.instantiate(cfg.detect_model)
# import pdb; pdb.set_trace()
# instantiate system
print_only(f"Instantiating system <{cfg.system._target_}>")
system: LightningModule = hydra.utils.instantiate(
cfg.system,
decouple_model=decouple_model,
detect_model=detect_model,
)
# instantiate callbacks
callbacks: List[Callback] = []
if cfg.get("early_stopping"):
print_only(f"Instantiating early_stopping <{cfg.early_stopping._target_}>")
callbacks.append(hydra.utils.instantiate(cfg.early_stopping))
if cfg.get("checkpoint"):
print_only(f"Instantiating checkpoint <{cfg.checkpoint._target_}>")
checkpoint: pl.callbacks.ModelCheckpoint = hydra.utils.instantiate(cfg.checkpoint)
callbacks.append(checkpoint)
# instantiate logger
print_only(f"Instantiating logger <{cfg.logger._target_}>")
os.makedirs(os.path.join(cfg.exp.dir, cfg.exp.name, "logs"), exist_ok=True)
logger = hydra.utils.instantiate(cfg.logger)
# instantiate trainer
print_only(f"Instantiating trainer <{cfg.trainer._target_}>")
trainer: Trainer = hydra.utils.instantiate(
cfg.trainer,
callbacks=callbacks,
logger=logger,
strategy=DDPStrategy(find_unused_parameters=True),
)
trainer.fit(system, datamodule=datamodule)
print_only("Training finished!")
best_k = {k: v.item() for k, v in checkpoint.best_k_models.items()}
with open(os.path.join(cfg.exp.dir, cfg.exp.name, "best_k_models.json"), "w") as f:
json.dump(best_k, f, indent=0)
import wandb
if wandb.run:
print_only("Closing wandb!")
wandb.finish()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--conf_dir",
default="local/conf.yml",
help="Full path to save best validation model",
)
args = parser.parse_args()
cfg = OmegaConf.load(args.conf_dir)
os.makedirs(os.path.join(cfg.exp.dir, cfg.exp.name), exist_ok=True)
# 保存配置到新的文件
OmegaConf.save(cfg, os.path.join(cfg.exp.dir, cfg.exp.name, "config.yaml"))
train(cfg)
|