File size: 2,179 Bytes
c0741ab |
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 |
import torch
import torch.nn as nn
import os
from model import MiniText
# -----------------------
# hiperparâmetros
# -----------------------
SEQ_LEN = 64
EPOCHS = 12000
LR = 1e-4
SAVE_EVERY = 2000 # salva checkpoint a cada X epochs
CHECKPOINT_PATH = "checkpoint.pt"
# -----------------------
# dataset
# -----------------------
with open("dataset.txt", "rb") as f:
data = torch.tensor(list(f.read()), dtype=torch.long)
# -----------------------
# model + optimizer
# -----------------------
model = MiniText()
optimizer = torch.optim.Adam(model.parameters(), lr=LR)
loss_fn = nn.CrossEntropyLoss()
start_epoch = 0
# -----------------------
# load checkpoint (se existir)
# -----------------------
if os.path.exists(CHECKPOINT_PATH):
print("Checkpoint encontrado, retomando treino...")
checkpoint = torch.load(CHECKPOINT_PATH)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
start_epoch = checkpoint["epoch"] + 1
else:
print("Nenhum checkpoint encontrado, treino do zero.")
# -----------------------
# batch sampler
# -----------------------
def get_batch():
idx = torch.randint(0, len(data) - SEQ_LEN - 1, (1,))
x = data[idx:idx + SEQ_LEN].unsqueeze(0)
y = data[idx + 1:idx + SEQ_LEN + 1].unsqueeze(0)
return x, y
# -----------------------
# training loop
# -----------------------
for epoch in range(start_epoch, EPOCHS):
x, y = get_batch()
logits, _ = model(x)
loss = loss_fn(logits.view(-1, 256), y.view(-1))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}/{EPOCHS} | Loss: {loss.item():.4f}")
# salvar checkpoint
if (epoch + 1) % SAVE_EVERY == 0:
torch.save({
"epoch": epoch,
"model": model.state_dict(),
"optimizer": optimizer.state_dict()
}, CHECKPOINT_PATH)
print("Checkpoint salvo.")
# -----------------------
# salvar modelo final
# -----------------------
torch.save(model.state_dict(), "minitext.pt")
print("Treino finalizado. Modelo salvo em minitext.pt")
|