(problem_type, boost_from_average)
| 1701 | ] |
| 1702 | ) |
| 1703 | def test_onnx_import(problem_type, boost_from_average): |
| 1704 | if (problem_type == 'binclass') and boost_from_average: |
| 1705 | pytest.xfail('CatBoost does not support importing classification onnx models with bias') |
| 1706 | |
| 1707 | if problem_type == 'binclass': |
| 1708 | loss_function = 'Logloss' |
| 1709 | train_path = TRAIN_FILE |
| 1710 | test_path = TEST_FILE |
| 1711 | cd_path = CD_FILE |
| 1712 | elif problem_type == 'multiclass': |
| 1713 | loss_function = 'MultiClass' |
| 1714 | train_path = CLOUDNESS_TRAIN_FILE |
| 1715 | test_path = CLOUDNESS_TEST_FILE |
| 1716 | cd_path = CLOUDNESS_CD_FILE |
| 1717 | elif problem_type == 'regression': |
| 1718 | loss_function = 'RMSE' |
| 1719 | train_path = TRAIN_FILE |
| 1720 | test_path = TEST_FILE |
| 1721 | cd_path = CD_FILE |
| 1722 | else: |
| 1723 | raise Exception('Unsupported problem_type: %s' % problem_type) |
| 1724 | |
| 1725 | train_pool = Pool(train_path, column_description=cd_path) |
| 1726 | test_pool = Pool(test_path, column_description=cd_path) |
| 1727 | |
| 1728 | model = CatBoost( |
| 1729 | { |
| 1730 | 'task_type': 'CPU', |
| 1731 | 'loss_function': loss_function, |
| 1732 | 'iterations': 5, |
| 1733 | 'depth': 4, |
| 1734 | 'ignored_features': train_pool.get_cat_feature_indices(), |
| 1735 | 'boost_from_average': boost_from_average |
| 1736 | } |
| 1737 | ) |
| 1738 | |
| 1739 | model.fit(train_pool) |
| 1740 | |
| 1741 | output_onnx_model_path = test_output_path(OUTPUT_ONNX_MODEL_PATH) |
| 1742 | model.save_model( |
| 1743 | output_onnx_model_path, |
| 1744 | format="onnx", |
| 1745 | export_parameters={ |
| 1746 | 'onnx_domain': 'ai.catboost', |
| 1747 | 'onnx_model_version': 1, |
| 1748 | 'onnx_doc_string': 'test model for problem_type %s' % problem_type, |
| 1749 | 'onnx_graph_name': 'CatBoostModel_for_%s' % problem_type |
| 1750 | } |
| 1751 | ) |
| 1752 | model.save_model(output_onnx_model_path, format="onnx") |
| 1753 | |
| 1754 | prediction_type = 'RawFormulaVal' if problem_type == 'regression' else 'Class' |
| 1755 | canon_pred = model.predict(test_pool, prediction_type=prediction_type) |
| 1756 | |
| 1757 | onnx_loaded_model = CatBoost( |
| 1758 | { |
| 1759 | 'task_type': 'CPU', |
| 1760 | 'loss_function': loss_function, |
nothing calls this directly
no test coverage detected