pass init config to initializer

KraFusion · Sep 12, 2017 · 1f13f10 · 1f13f10
1 parent d5f9485
commit 1f13f10
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/.../src/main/java/org/deeplearning4j/nn/modelimport/keras/utils/KerasInitilizationUtils.java b/.../src/main/java/org/deeplearning4j/nn/modelimport/keras/utils/KerasInitilizationUtils.java
@@ -43,7 +43,9 @@ public class KerasInitilizationUtils {
      * @return DL4J weight initialization enum
      * @see WeightInit
      */
-    public static Pair<WeightInit, Distribution> mapWeightInitialization(String kerasInit, KerasLayerConfiguration conf)
+    public static Pair<WeightInit, Distribution> mapWeightInitialization(String kerasInit,
+                                                                         KerasLayerConfiguration conf,
+                                                                         Map<String, Object> initConfig)
             throws UnsupportedKerasConfigurationException {
 
         WeightInit init = WeightInit.XAVIER;
@@ -133,17 +135,19 @@ public static WeightInit getWeightInitFromConfig(Map<String, Object> layerConfig
         if (!innerConfig.containsKey(initField))
             throw new InvalidKerasConfigurationException("Keras layer is missing " + initField + " field");
         String kerasInit = "glorot_normal";
-        if (kerasMajorVersion != 2)
+        Map<String, Object> initMap;
+        if (kerasMajorVersion != 2) {
             kerasInit = (String) innerConfig.get(initField);
-        else {
-            HashMap initMap = (HashMap) innerConfig.get(initField);
+            initMap = innerConfig;
+        } else {
+            initMap = (HashMap) innerConfig.get(initField);
             if (initMap.containsKey("class_name")) {
                 kerasInit = (String) initMap.get("class_name");
             }
         }
         Pair<WeightInit, Distribution> init;
         try {
-            init = mapWeightInitialization(kerasInit, conf);
+            init = mapWeightInitialization(kerasInit, conf, initMap);
         } catch (UnsupportedKerasConfigurationException e) {
             if (enforceTrainingConfig)
                 throw e;

diff --git a/...t/java/org/deeplearning4j/nn/modelimport/keras/configurations/KerasInitilizationTest.java b/...t/java/org/deeplearning4j/nn/modelimport/keras/configurations/KerasInitilizationTest.java
@@ -58,7 +58,7 @@ private WeightInit[] dl4jInitializers() {
         return new WeightInit[] {
                 WeightInit.XAVIER,
                 WeightInit.XAVIER_UNIFORM,
-                WeightInit.NORMAL,
+                WeightInit.LECUN_NORMAL,
                 WeightInit.UNIFORM,
                 WeightInit.RELU,
                 WeightInit.RELU_UNIFORM,