Spaces:
Sleeping
Sleeping
File size: 1,921 Bytes
86e31ba 06aa4a6 86e31ba |
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 |
import torch
from torch.utils.data import DataLoader
from transformers import AdamW, get_linear_schedule_with_warmup
# from utils.dataset import BDTtsDataset
from inference import tts # reuse your model
training_config = {
"learning_rate": 1e-4,
"batch_size": 16,
"warmup_steps": 1000,
"gradient_accumulation_steps": 4,
"mixed_precision": True,
"save_strategy": "steps",
"save_steps": 500,
"eval_steps": 100,
"num_epochs": 5
}
def train():
dataset = BDTtsDataset("./data/train")
dataloader = DataLoader(dataset, batch_size=training_config["batch_size"], shuffle=True)
optimizer = AdamW(tts.model.parameters(), lr=training_config["learning_rate"])
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=training_config["warmup_steps"],
num_training_steps=len(dataloader) * training_config["num_epochs"]
)
scaler = torch.cuda.amp.GradScaler() if training_config["mixed_precision"] else None
step = 0
for epoch in range(training_config["num_epochs"]):
for batch in dataloader:
inputs, targets = batch
optimizer.zero_grad()
with torch.cuda.amp.autocast(enabled=scaler is not None):
outputs = tts.model(inputs)
loss = outputs.loss if hasattr(outputs, "loss") else torch.nn.functional.mse_loss(outputs, targets)
if scaler:
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
optimizer.step()
scheduler.step()
step += 1
if step % training_config["save_steps"] == 0:
torch.save(tts.model.state_dict(), f"checkpoints/model_step{step}.pth")
print(f"Saved checkpoint at step {step}")
if __name__ == "__main__":
train()
|