Spaces:
Running on Zero
Running on Zero
File size: 1,846 Bytes
cc348e2 | 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 | """
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)
|