Skip to content

Commit 8b8a6d2

Browse files
author
Omede Firouz
committed
[SPARK-7022][PySpark][ML] Add ML.Tuning.ParamGridBuilder to PySpark
1 parent 3ae37b9 commit 8b8a6d2

File tree

2 files changed

+85
-0
lines changed

2 files changed

+85
-0
lines changed

python/pyspark/ml/tuning.py

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one or more
3+
# contributor license agreements. See the NOTICE file distributed with
4+
# this work for additional information regarding copyright ownership.
5+
# The ASF licenses this file to You under the Apache License, Version 2.0
6+
# (the "License"); you may not use this file except in compliance with
7+
# the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
18+
__all__ = ['ParamGridBuilder']
19+
20+
class ParamGridBuilder(object):
21+
"""
22+
Builder for a param grid used in grid search-based model selection.
23+
"""
24+
25+
def __init__(self):
26+
self._param_grid = {}
27+
28+
def addGrid(self, param, values):
29+
"""
30+
Sets the given parameters in this grid to fixed values.
31+
"""
32+
self._param_grid[param] = values
33+
34+
def baseOn(self, *args):
35+
"""
36+
Sets the given parameters in this grid to fixed values.
37+
Accepts either a parameter dictionary or a list of (parameter, value) pairs.
38+
"""
39+
if isinstance(args[0], dict):
40+
self.baseOn(*args[0].items())
41+
else:
42+
for (param, value) in args:
43+
self.addGrid(param, [value])
44+
45+
def build(self):
46+
"""
47+
Builds and returns all combinations of parameters specified
48+
by the param grid.
49+
"""
50+
param_maps = [{}]
51+
for (param, values) in self._param_grid.items():
52+
new_param_maps = []
53+
for value in values:
54+
for old_map in param_maps:
55+
copied_map = old_map.copy()
56+
copied_map[param] = value
57+
new_param_maps.append(copied_map)
58+
param_maps = new_param_maps
59+
60+
return param_maps
61+
62+
63+
if __name__ == "__main__":
64+
grid_test = ParamGridBuilder()
65+
from classification import LogisticRegression
66+
lr = LogisticRegression()
67+
grid_test.addGrid(lr.regParam, [1.0, 2.0, 3.0])
68+
grid_test.addGrid(lr.maxIter, [1, 5])
69+
grid_test.addGrid(lr.featuresCol, ['f'])
70+
grid_test.baseOn({lr.labelCol: 'l'})
71+
grid_test.baseOn([lr.predictionCol, 'p'])
72+
grid = grid_test.build()
73+
expected = [
74+
{lr.regParam: 1.0, lr.featuresCol: 'f', lr.maxIter: 1, lr.labelCol: 'l', lr.predictionCol: 'p'},
75+
{lr.regParam: 2.0, lr.featuresCol: 'f', lr.maxIter: 1, lr.labelCol: 'l', lr.predictionCol: 'p'},
76+
{lr.regParam: 3.0, lr.featuresCol: 'f', lr.maxIter: 1, lr.labelCol: 'l', lr.predictionCol: 'p'},
77+
{lr.regParam: 1.0, lr.featuresCol: 'f', lr.maxIter: 5, lr.labelCol: 'l', lr.predictionCol: 'p'},
78+
{lr.regParam: 2.0, lr.featuresCol: 'f', lr.maxIter: 5, lr.labelCol: 'l', lr.predictionCol: 'p'},
79+
{lr.regParam: 3.0, lr.featuresCol: 'f', lr.maxIter: 5, lr.labelCol: 'l', lr.predictionCol: 'p'}
80+
]
81+
82+
for a, b in zip(grid, expected):
83+
if a != b:
84+
exit(-1)

python/run-tests

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ function run_ml_tests() {
9494
echo "Run ml tests ..."
9595
run_test "pyspark/ml/feature.py"
9696
run_test "pyspark/ml/classification.py"
97+
run_test "pyspark/ml/tuning.py"
9798
run_test "pyspark/ml/tests.py"
9899
}
99100

0 commit comments

Comments
 (0)