Skip to content

count completed training steps #58

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
borawhocodess wants to merge 1 commit into
base: abstraction
Choose a base branch
from
Open

count completed training steps #58

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 12 additions & 10 deletions tfmplayground/training/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,7 @@ def train(
"""
trains model on prior batches for given epochs

steps with non-finite values are skipped, and an epoch where every step is
skipped raises
batches with non-finite values are retried, and 100 consecutive skips raise

Parameters
----------
Expand All @@ -48,7 +47,7 @@ def train(
batch_size : int
number of tables per batch
steps_per_epoch : int
number of batches that make up one epoch
number of optimizer updates per epoch
lr : float
learning rate for optimizer
grad_clip : float
Expand All @@ -65,6 +64,8 @@ def train(
TabularFoundationModel
model after last epoch
"""
if steps_per_epoch < 1:
raise ValueError("steps_per_epoch must be at least 1")
if callbacks is None:
callbacks = []
if not device:
Expand All @@ -83,7 +84,11 @@ def train(
optimizer.train()
total_loss = 0.0
num_valid = 0
for _ in range(steps_per_epoch):
consecutive_skips = 0
while num_valid < steps_per_epoch:
if consecutive_skips >= 100:
raise RuntimeError(f"100 consecutive batches in epoch {epoch} had non-finite values")
consecutive_skips += 1
x_train, y_train, x_test, y_test = next(batches)
x_train = x_train.to(device)
y_train = y_train.to(device)
Expand Down Expand Up @@ -117,15 +122,12 @@ def train(

loss = losses.mean()
loss.backward()
total_loss += loss.cpu().detach().item()
num_valid += 1

torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
optimizer.step()
optimizer.zero_grad()

if num_valid == 0:
raise RuntimeError(f"all {steps_per_epoch} steps in epoch {epoch} had non-finite values")
total_loss += loss.cpu().detach().item()
num_valid += 1
consecutive_skips = 0

end_time = time.time()
mean_loss = total_loss / num_valid
Expand Down