LightGBM pipeline serialization fix (dotnet#251)

Dmitry-A · Mar 3, 2019 · b3fd4dc · b3fd4dc
1 parent 8fd2aa8
commit b3fd4dc
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 6 deletions.
diff --git a/src/Microsoft.ML.Auto/TrainerExtensions/TrainerExtensionUtil.cs b/src/Microsoft.ML.Auto/TrainerExtensions/TrainerExtensionUtil.cs
@@ -131,7 +131,7 @@ private static IDictionary<string, object> BuildLightGbmPipelineNodeProps(IEnume
             string labelColumn, string weightColumn)
         {
             Dictionary<string, object> props = null;
-            if (sweepParams == null)
+            if (sweepParams == null || !sweepParams.Any())
             {
                 props = new Dictionary<string, object>();
             }
@@ -185,11 +185,19 @@ public static ColumnInformation BuildColumnInfo(IDictionary<string, object> prop
 
         private static ParameterSet BuildLightGbmParameterSet(IDictionary<string, object> props)
         {
-            var parentProps = props.Where(p => p.Key != LightGbmTreeBoosterPropName);
-            var treeProps = ((CustomProperty)props[LightGbmTreeBoosterPropName]).Properties;
-            var allProps = parentProps.Union(treeProps);
-            var paramVals = allProps.Select(p => new StringParameterValue(p.Key, p.Value.ToString()));
-            return new ParameterSet(paramVals);
+            IEnumerable<IParameterValue> parameters;
+            if (props == null || !props.Any())
+            {
+                parameters = new List<IParameterValue>();
+            }
+            else
+            {
+                var parentProps = props.Where(p => p.Key != LightGbmTreeBoosterPropName);
+                var treeProps = ((CustomProperty)props[LightGbmTreeBoosterPropName]).Properties;
+                var allProps = parentProps.Union(treeProps);
+                parameters = allProps.Select(p => new StringParameterValue(p.Key, p.Value.ToString()));
+            }
+            return new ParameterSet(parameters);
         }
 
         private static void SetValue(FieldInfo fi, IComparable value, object obj, Type propertyType)

diff --git a/src/Test/TrainerExtensionsTests.cs b/src/Test/TrainerExtensionsTests.cs
@@ -111,6 +111,28 @@ public void BuildSdcaPipelineNode()
             Util.AssertObjectMatchesJson(expectedJson, pipelineNode);
         }
 
+        [TestMethod]
+        public void BuildLightGbmPipelineNodeDefaultParams()
+        {
+            var pipelineNode = new LightGbmBinaryExtension().CreatePipelineNode(
+                new List<SweepableParam>(), 
+                new ColumnInformation());
+            var expectedJson = @"{
+  ""Name"": ""LightGbmBinary"",
+  ""NodeType"": ""Trainer"",
+  ""InColumns"": [
+    ""Features""
+  ],
+  ""OutColumns"": [
+    ""Score""
+  ],
+  ""Properties"": {
+    ""LabelColumn"": ""Label""
+  }
+}";
+            Util.AssertObjectMatchesJson(expectedJson, pipelineNode);
+        }
+
         [TestMethod]
         public void BuildPipelineNodeWithCustomColumns()
         {