free dataset (microsoft#427)

Alnusjaponica · Apr 18, 2017 · 68e2530 · 68e2530
1 parent 6318b79
commit 68e2530
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/python-package/lightgbm/basic.py b/python-package/lightgbm/basic.py
@@ -1267,6 +1267,10 @@ def __setstate__(self, state):
             state['handle'] = handle
         self.__dict__.update(state)
 
+    def free_dataset(self):
+        self.__dict__.pop('train_set', None)
+        self.__dict__.pop('valid_sets', None)
+
     def set_train_data_name(self, name):
         self.__train_data_name = name
 

diff --git a/python-package/lightgbm/sklearn.py b/python-package/lightgbm/sklearn.py
@@ -418,6 +418,10 @@ def get_meta_data(collection, i):
         if early_stopping_rounds is not None:
             self.best_iteration = self._Booster.best_iteration
         self.best_score = self._Booster.best_score
+
+        # free dataset
+        self.booster_.free_dataset()
+        del train_set, valid_sets
         return self
 
     def predict(self, X, raw_score=False, num_iteration=0):