diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 78689f52..9744a9d3 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -33,6 +33,7 @@ class Trainer(TrainerIO): log_save_interval=1, add_log_row_interval=1, lr_scheduler_milestones=None, use_amp=False, + check_grad_nans=False, amp_level='O2', nb_sanity_val_steps=5): @@ -60,6 +61,7 @@ class Trainer(TrainerIO): self.lr_scheduler_milestones = [] if lr_scheduler_milestones is None else [int(x.strip()) for x in lr_scheduler_milestones.split(',')] self.lr_schedulers = [] self.amp_level = amp_level + self.check_grad_nans = check_grad_nans # training state self.optimizers = None @@ -373,11 +375,13 @@ class Trainer(TrainerIO): for optimizer in self.optimizers: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() - for param in self.model.parameters(): - print(param.grad.float().sum()) else: loss.backward() + if self.check_grad_nans: + for param in self.model.parameters(): + print(param.grad.float().sum()) + self.batch_loss_value += loss.item() # gradient update with accumulated gradients diff --git a/pytorch_lightning/utils/arg_parse.py b/pytorch_lightning/utils/arg_parse.py index e72eb4e6..4d05e297 100644 --- a/pytorch_lightning/utils/arg_parse.py +++ b/pytorch_lightning/utils/arg_parse.py @@ -50,7 +50,8 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None parser.add_argument('--single_run_gpu', dest='single_run_gpu', action='store_true') parser.add_argument('--disable_cuda', dest='disable_cuda', action='store_true') parser.add_argument('--default_tensor_type', default='torch.cuda.FloatTensor', type=str) - parser.add_argument('--use_amp', dest='use_amp', action='store_true') + parser.add_argument('--use_amp', dest='use_amp', action='store_true')a= + parser.add_argument('--check_grad_nans', dest='check_grad_nans', action='store_true')a= parser.add_argument('--amp_level', default='O2',type=str)