Spaces:
Running on Zero
Running on Zero
File size: 2,343 Bytes
98d1e65 972b71a 98d1e65 972b71a 98d1e65 972b71a 98d1e65 | 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 104 105 106 107 108 109 110 | import torch
import torch.nn as nn
from src.infra.dataset.autoregressive import AutoRegressiveDataset
from src.infra.dataset.util import get_batch
# Teach 'em
def train(
model,
train_dataset: AutoRegressiveDataset,
val_dataset: AutoRegressiveDataset,
training_config,
):
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(
model.parameters(),
lr=training_config.learning_rate,
)
train_history = []
val_history = []
steps = []
model.train()
for step in range(training_config.num_steps):
if step % training_config.eval_interval == 0:
train_loss = evaluate(
model=model,
dataset=train_dataset,
training_config=training_config,
)
val_loss = evaluate(
model=model,
dataset=val_dataset,
training_config=training_config,
)
steps.append(step)
train_history.append(train_loss)
val_history.append(val_loss)
print(
f"step {step:4d} | "
f"train {train_loss:.4f} | "
f"val {val_loss:.4f}"
)
inputs, targets = get_batch(
dataset=train_dataset,
batch_size=training_config.batch_size,
)
logits = model(inputs)
loss = loss_function(
logits.view(-1, logits.size(-1)),
targets.view(-1),
)
optimizer.zero_grad()
loss.backward()
optimizer.step()
#return model
return model, steps, train_history, val_history
@torch.no_grad() #do not calculate gradients
def evaluate(
model,
dataset: AutoRegressiveDataset,
training_config,
):
model.eval()
loss_function = nn.CrossEntropyLoss()
losses = []
for _ in range(training_config.eval_batches):
inputs, targets = get_batch(
dataset=dataset,
batch_size=training_config.batch_size,
)
logits = model(inputs)
loss = loss_function(
logits.view(-1, logits.size(-1)),
targets.view(-1),
)
losses.append(loss.item())
model.train()
return sum(losses) / len(losses)
|