| 7136 | |
| 7137 | |
| 7138 | def test_shrink(): |
| 7139 | train_pool = Pool(TRAIN_FILE, column_description=CD_FILE) |
| 7140 | test_pool = Pool(TEST_FILE, column_description=CD_FILE) |
| 7141 | args = { |
| 7142 | 'iterations': 30, |
| 7143 | 'loss_function': 'Logloss', |
| 7144 | 'use_best_model': False, |
| 7145 | 'learning_rate': 0.3 |
| 7146 | } |
| 7147 | model = CatBoostClassifier(**args) |
| 7148 | args['iterations'] = 9 |
| 7149 | model2 = CatBoostClassifier(**args) |
| 7150 | |
| 7151 | model.fit(train_pool, eval_set=test_pool) |
| 7152 | model2.fit(train_pool, eval_set=test_pool) |
| 7153 | assert model.tree_count_ == 30 |
| 7154 | model.shrink(9) |
| 7155 | assert model.tree_count_ == 9 |
| 7156 | pred1 = model.predict(test_pool) |
| 7157 | pred2 = model2.predict(test_pool) |
| 7158 | assert np.all(pred1 == pred2) |
| 7159 | model.shrink(8, ntree_start=1) |
| 7160 | assert model.tree_count_ == 7 |
| 7161 | |
| 7162 | |
| 7163 | def test_shrink_with_bias(): |