Spaces:
Running on Zero
Running on Zero
File size: 1,766 Bytes
1717c0c | 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 | import torch
import torch.nn as nn
from torchgen import model
from src.infra.dataset.answering import QuestionAnsweringDataset
from torch.utils.data import DataLoader
from src.task.model import TinyAnsweringModel
# Teach 'em
def train(
model: TinyAnsweringModel,
train_dataset: QuestionAnsweringDataset,
val_dataset: QuestionAnsweringDataset,
training_config,
):
train_loader = DataLoader(
train_dataset,
batch_size=training_config.batch_size,
shuffle=True,
)
num_batches = len(train_loader)
print(f"Training examples: {len(train_dataset)}")
print(f"Batch size: {training_config.batch_size}")
print(f"Batches per epoch: {num_batches}")
print(f"Total batches: {num_batches * training_config.num_steps}")
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(
model.parameters(),
lr=training_config.learning_rate,
)
model.train()
for epoch in range(training_config.num_steps):
print(f"Epoch {epoch+1}/{training_config.num_steps}")
for tokens, padding_mask, start_position, end_position in train_loader:
start_logits, end_logits = model(tokens,padding_mask)
start_position = start_position.long()
end_position = end_position.long() -1 #bounds bug
start_loss = loss_function(start_logits,start_position)
end_loss = loss_function(end_logits,end_position)
loss = start_loss + end_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
#poor-man's eval
if epoch % training_config.eval_interval == 0:
print(f"\tLoss : {loss}")
return model
|