From 7d9ffe6d4ecbadd688eb3ebebb6c6c370281820f Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 29 Aug 2019 14:29:07 +0200 Subject: [PATCH] weight_decay fix --- train.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/train.py b/train.py index 938e19e5..127adad2 100644 --- a/train.py +++ b/train.py @@ -258,9 +258,6 @@ def train(): # Compute loss loss, loss_items = compute_loss(pred, targets, model) - if torch.isnan(loss): - print('WARNING: nan loss detected, ending training', loss_items) - return results # Scale loss by nominal batch_size of 64 loss *= batch_size / 64 @@ -282,8 +279,14 @@ def train(): mem = torch.cuda.memory_cached() / 1E9 if torch.cuda.is_available() else 0 # (GB) s = ('%10s' * 2 + '%10.3g' * 6) % ( '%g/%g' % (epoch, epochs - 1), '%.3gG' % mem, *mloss, len(targets), img_size) - pbar.set_description(s) # end batch ----------------------------------------------------------------------- + pbar.set_description(s) + if torch.isnan(loss): + print('WARNING: nan loss detected, ending training', loss_items) + return results + # end batch ------------------------------------------------------------------------------------------------ + + # Process epoch results final_epoch = epoch + 1 == epochs if opt.prebias: print_model_biases(model) @@ -342,7 +345,9 @@ def train(): torch.save(chkpt, wdir + 'backup%g.pt' % epoch) # Delete checkpoint - del chkpt # end epoch ------------------------------------------------------------------------------------- + del chkpt + + # end epoch ---------------------------------------------------------------------------------------------------- # Report time plot_results() # save as results.png