Fix issue with regularization parameters not being passed.

Author: otiliastr

neural_structured_learning/research/gam/trainer/trainer_classification.py

@@ -209,15 +209,18 @@ def __init__(self,
 
     # Create variables and predictions.
     with tf.variable_scope('predictions'):
-      encoding, variables, reg_params = self.model.get_encoding_and_params(
-          inputs=input_features, is_train=is_train)
-      self.variables = variables
-      self.reg_params = reg_params
-      predictions, variables, reg_params = (
+      encoding, variables_enc, reg_params_enc = (
+          self.model.get_encoding_and_params(
+            inputs=input_features,
+            is_train=is_train))
+      self.variables = variables_enc
+      self.reg_params = reg_params_enc
+      predictions, variables_pred, reg_params_pred = (
           self.model.get_predictions_and_params(
-              encoding=encoding, is_train=is_train))
-      self.variables.update(variables)
-      self.reg_params.update(reg_params)
+              encoding=encoding,
+              is_train=is_train))
+      self.variables.update(variables_pred)
+      self.reg_params.update(reg_params_pred)
       normalized_predictions = self.model.normalize_predictions(predictions)
       predictions_var_scope = tf.get_variable_scope()
 
@@ -262,7 +265,7 @@ def __init__(self,
       # Weight decay loss.
       loss_reg = 0.0
       if weight_decay_var is not None:
-        for var in reg_params.values():
+        for var in self.reg_params.values():
           loss_reg += weight_decay_var * tf.nn.l2_loss(var)
 
       # Adversarial loss, in case we want to add VAT on top of GAM.
@@ -351,7 +354,7 @@ def __init__(self,
     if isinstance(weight_decay_var, tf.Variable):
       self.vars_to_save.append(weight_decay_var)
     if self.warm_start:
-      self.vars_to_save.extend([v for v in variables])
+      self.vars_to_save.extend([v for v in self.variables])
 
     # More variables to be initialized after the session is created.
     self.is_initialized = False
@@ -366,7 +369,6 @@ def __init__(self,
     self.weight_decay_update = weight_decay_update
     self.iter_cls_total = iter_cls_total
     self.iter_cls_total_update = iter_cls_total_update
-    self.variables = variables
     self.accuracy = accuracy
     self.train_op = train_op
     self.loss_op = loss_op