From 218f0a5b4a42b610959b062d2bdedb82c78f7b6e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Wed, 4 Dec 2019 13:32:47 +0100 Subject: [PATCH] inspect training_step for opt_idx (#573) --- pytorch_lightning/trainer/model_hooks_mixin.py | 6 ++++++ pytorch_lightning/trainer/train_loop_mixin.py | 10 +++++++++- 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/pytorch_lightning/trainer/model_hooks_mixin.py b/pytorch_lightning/trainer/model_hooks_mixin.py index 23c82485..442c5209 100644 --- a/pytorch_lightning/trainer/model_hooks_mixin.py +++ b/pytorch_lightning/trainer/model_hooks_mixin.py @@ -1,3 +1,4 @@ +import inspect from pytorch_lightning.core.lightning import LightningModule @@ -15,3 +16,8 @@ class TrainerModelHooksMixin(object): # when code pointers are different, it was overriden is_overriden = getattr(model, f_name).__code__ is not getattr(super_object, f_name).__code__ return is_overriden + + def has_arg(self, f_name, arg_name): + model = self.get_model() + f_op = getattr(model, f_name, None) + return arg_name in inspect.signature(f_op).parameters diff --git a/pytorch_lightning/trainer/train_loop_mixin.py b/pytorch_lightning/trainer/train_loop_mixin.py index 0bff3d59..49cd3d70 100644 --- a/pytorch_lightning/trainer/train_loop_mixin.py +++ b/pytorch_lightning/trainer/train_loop_mixin.py @@ -149,6 +149,7 @@ When this flag is enabled each batch is split into sequences of size truncated_b """ +import inspect import numpy as np from pytorch_lightning.utilities.debugging import MisconfigurationException @@ -434,8 +435,15 @@ class TrainerTrainLoopMixin(object): # --------------- # enable not needing to add opt_idx to training_step args = [batch, batch_idx] + if len(self.optimizers) > 1: - args.append(opt_idx) + if self.has_arg('training_step', 'optimizer_idx'): + args.append(opt_idx) + else: + raise ValueError( + f'Your LightningModule defines {len(self.optimizers)} optimizers but ' + f'training_step is missing the "optimizer_idx" argument.' + ) # pass hiddens if using tbptt if self.truncated_bptt_steps is not None: