Spaces:
Running on Zero
Running on Zero
| """ | |
| This is basically an integration test | |
| """ | |
| import random | |
| import torch | |
| from TrainingInterfaces.Text_to_Spectrogram.FastSpeech2.FastSpeech2 import FastSpeech2 | |
| from TrainingInterfaces.Text_to_Spectrogram.FastSpeech2.fastspeech2_train_loop import train_loop | |
| from Utility.corpus_preparation import prepare_fastspeech_corpus | |
| from Utility.path_to_transcript_dicts import * | |
| def run(gpu_id, resume_checkpoint, finetune, model_dir, resume): | |
| if gpu_id == "cpu": | |
| os.environ["CUDA_VISIBLE_DEVICES"] = "" | |
| device = torch.device("cpu") | |
| else: | |
| os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" | |
| os.environ["CUDA_VISIBLE_DEVICES"] = f"{gpu_id}" | |
| device = torch.device("cuda") | |
| torch.manual_seed(131714) | |
| random.seed(131714) | |
| torch.random.manual_seed(131714) | |
| print("Preparing") | |
| if model_dir is not None: | |
| save_dir = model_dir | |
| else: | |
| save_dir = os.path.join("Models", "FastSpeech2_IntegrationTest") | |
| os.makedirs(save_dir, exist_ok=True) | |
| train_set = prepare_fastspeech_corpus(transcript_dict=build_path_to_transcript_dict_integration_test(), | |
| corpus_dir=os.path.join("Corpora", "IntegrationTest"), | |
| lang="en", | |
| save_imgs=True) | |
| model = FastSpeech2(lang_embs=None, utt_embed_dim=None) | |
| print("Training model") | |
| train_loop(net=model, | |
| train_dataset=train_set, | |
| device=device, | |
| save_directory=save_dir, | |
| steps=5000, | |
| batch_size=32, | |
| lang="en", | |
| lr=0.001, | |
| epochs_per_save=1, | |
| warmup_steps=4000, | |
| path_to_checkpoint=resume_checkpoint, | |
| fine_tune=finetune, | |
| resume=resume) | |